diff --git a/.gitignore b/.gitignore index c2592bc5..f124cc14 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,4 @@ out/ screenshots/ headless_output/ www/assets/ +.web-asset-cache diff --git a/Cargo.lock b/Cargo.lock index 3d0eda24..1aae0ddc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -180,6 +180,28 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "alsa" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed7572b7ba83a31e20d1b48970ee402d2e3e0537dcfe0a3ff4d6eb7508617d43" +dependencies = [ + "alsa-sys", + "bitflags 2.9.4", + "cfg-if", + "libc", +] + +[[package]] +name = "alsa-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db8fee663d06c4e303404ef5f40488a53e062f89ba8bfed81f42325aafad1527" +dependencies = [ + "libc", + "pkg-config", +] + [[package]] name = "android-activity" version = "0.6.0" @@ -194,7 +216,7 @@ dependencies = [ "jni-sys", "libc", "log", - "ndk", + "ndk 0.9.0", "ndk-context", "ndk-sys 0.6.0+11769913", "num_enum 0.7.4", @@ -470,6 +492,26 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "basis-universal" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "555fb05709f4e12fa2f6b93a480facf167eb0ecb2558ba41f610f588e77cbd14" +dependencies = [ + "basis-universal-sys", + "bitflags 1.3.2", + "lazy_static", +] + +[[package]] +name = "basis-universal-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd9bde5e9547958fb0e77d79fc7879edcf91d5e0c8e372ef8959916cf35e8506" +dependencies = [ + "cc", +] + [[package]] name = "bevy" version = "0.14.2" @@ -748,6 +790,23 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "bevy_audio" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83620c82f281848c02ed4b65133a0364512b4eca2b39cd21a171e50e2986d89" +dependencies = [ + "bevy_app 0.17.2", + "bevy_asset 0.17.2", + "bevy_ecs 0.17.2", + "bevy_math 0.17.2", + "bevy_reflect 0.17.2", + "bevy_transform 0.17.2", + "coreaudio-sys", + "rodio", + "tracing", +] + [[package]] name = "bevy_camera" version = "0.17.2" @@ -1114,6 +1173,7 @@ dependencies = [ "bevy_file_asset", "bevy_interleave", "bevy_panorbit_camera", + "bevy_pbr 0.17.2", "bevy_transform_gizmo", "bincode2", "byte-unit", @@ -1129,6 +1189,7 @@ dependencies = [ "kd-tree", "noise", "ply-rs", + "png 0.17.16", "pollster", "rand 0.8.5", "rayon", @@ -1214,6 +1275,40 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "bevy_gltf" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d67e954b20551818f7cdb33f169ab4db64506ada66eb4d60d3cb8861103411" +dependencies = [ + "base64 0.22.1", + "bevy_app 0.17.2", + "bevy_asset 0.17.2", + "bevy_camera", + "bevy_color 0.17.2", + "bevy_ecs 0.17.2", + "bevy_image", + "bevy_light", + "bevy_math 0.17.2", + "bevy_mesh", + "bevy_pbr 0.17.2", + "bevy_platform", + "bevy_reflect 0.17.2", + "bevy_render 0.17.2", + "bevy_scene 0.17.2", + "bevy_tasks 0.17.2", + "bevy_transform 0.17.2", + "fixedbitset 0.5.7", + "gltf", + "itertools 0.14.0", + "percent-encoding", + "serde", + "serde_json", + "smallvec", + "thiserror 2.0.17", + "tracing", +] + [[package]] name = "bevy_hierarchy" version = "0.14.2" @@ -1234,6 +1329,7 @@ version = "0.17.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "168de8239b2aedd2eeef9f76ae1909b2fdf859b11dcdb4d4d01b93f5f2c771be" dependencies = [ + "basis-universal", "bevy_app 0.17.2", "bevy_asset 0.17.2", "bevy_color 0.17.2", @@ -1248,6 +1344,7 @@ dependencies = [ "guillotiere", "half", "image", + "ktx2", "rectangle-pack", "ruzstd", "serde", @@ -1392,6 +1489,7 @@ dependencies = [ "bevy_anti_alias", "bevy_app 0.17.2", "bevy_asset 0.17.2", + "bevy_audio", "bevy_camera", "bevy_color 0.17.2", "bevy_core_pipeline 0.17.2", @@ -1399,6 +1497,7 @@ dependencies = [ "bevy_diagnostic 0.17.2", "bevy_ecs 0.17.2", "bevy_gizmos 0.17.2", + "bevy_gltf", "bevy_image", "bevy_input 0.17.2", "bevy_input_focus", @@ -1408,6 +1507,7 @@ dependencies = [ "bevy_mesh", "bevy_pbr 0.17.2", "bevy_platform", + "bevy_post_process", "bevy_ptr 0.17.2", "bevy_reflect 0.17.2", "bevy_render 0.17.2", @@ -1807,6 +1907,36 @@ dependencies = [ "web-time", ] +[[package]] +name = "bevy_post_process" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26ee8ab6043f8bbe43e9c16bbdde0c5e7289b99e62cd8aad1a2a4166a7f2bce6" +dependencies = [ + "bevy_app 0.17.2", + "bevy_asset 0.17.2", + "bevy_camera", + "bevy_color 0.17.2", + "bevy_core_pipeline 0.17.2", + "bevy_derive 0.17.2", + "bevy_ecs 0.17.2", + "bevy_image", + "bevy_math 0.17.2", + "bevy_platform", + "bevy_reflect 0.17.2", + "bevy_render 0.17.2", + "bevy_shader", + "bevy_transform 0.17.2", + "bevy_utils 0.17.2", + "bevy_window 0.17.2", + "bitflags 2.9.4", + "nonmax", + "radsort", + "smallvec", + "thiserror 2.0.17", + "tracing", +] + [[package]] name = "bevy_ptr" version = "0.14.2" @@ -2480,6 +2610,24 @@ dependencies = [ "serde", ] +[[package]] +name = "bindgen" +version = "0.72.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" +dependencies = [ + "bitflags 2.9.4", + "cexpr", + "clang-sys", + "itertools 0.10.5", + "proc-macro2", + "quote", + "regex", + "rustc-hash 2.1.1", + "shlex", + "syn 2.0.106", +] + [[package]] name = "bit-set" version = "0.5.3" @@ -2774,6 +2922,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "1.0.3" @@ -2819,6 +2976,17 @@ dependencies = [ "half", ] +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading 0.8.9", +] + [[package]] name = "clap" version = "4.5.48" @@ -3062,6 +3230,26 @@ dependencies = [ "libc", ] +[[package]] +name = "coreaudio-rs" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "321077172d79c662f64f5071a03120748d5bb652f5231570141be24cfcd2bace" +dependencies = [ + "bitflags 1.3.2", + "core-foundation-sys", + "coreaudio-sys", +] + +[[package]] +name = "coreaudio-sys" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ceec7a6067e62d6f931a2baf6f3a751f4a892595bcec1461a3c94ef9949864b6" +dependencies = [ + "bindgen", +] + [[package]] name = "cosmic-text" version = "0.14.2" @@ -3072,7 +3260,7 @@ dependencies = [ "fontdb", "log", "rangemap", - "rustc-hash", + "rustc-hash 1.1.0", "rustybuzz", "self_cell", "smol_str", @@ -3085,6 +3273,29 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cpal" +version = "0.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "873dab07c8f743075e57f524c583985fbaf745602acbe916a01539364369a779" +dependencies = [ + "alsa", + "core-foundation-sys", + "coreaudio-rs", + "dasp_sample", + "jni", + "js-sys", + "libc", + "mach2", + "ndk 0.8.0", + "ndk-context", + "oboe", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "windows 0.54.0", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -3229,6 +3440,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "dasp_sample" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c87e182de0887fd5361989c677c4e8f5000cd9491d6d563161a8f3a5519fc7f" + [[package]] name = "data-encoding" version = "2.9.0" @@ -3953,6 +4170,42 @@ dependencies = [ "web-sys", ] +[[package]] +name = "gltf" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3ce1918195723ce6ac74e80542c5a96a40c2b26162c1957a5cd70799b8cacf7" +dependencies = [ + "byteorder", + "gltf-json", + "lazy_static", + "serde_json", +] + +[[package]] +name = "gltf-derive" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14070e711538afba5d6c807edb74bcb84e5dbb9211a3bf5dea0dfab5b24f4c51" +dependencies = [ + "inflections", + "proc-macro2", + "quote", + "syn 2.0.106", +] + +[[package]] +name = "gltf-json" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6176f9d60a7eab0a877e8e96548605dedbde9190a7ae1e80bbcc1c9af03ab14" +dependencies = [ + "gltf-derive", + "serde", + "serde_derive", + "serde_json", +] + [[package]] name = "glutin_wgl_sys" version = "0.5.0" @@ -4310,8 +4563,10 @@ dependencies = [ "byteorder-lite", "moxcms", "num-traits", - "png", + "png 0.18.0", "tiff", + "zune-core", + "zune-jpeg", ] [[package]] @@ -4335,6 +4590,12 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inflections" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a257582fdcde896fd96463bf2d40eefea0580021c0712a0e2b028b60b47a837a" + [[package]] name = "inventory" version = "0.3.21" @@ -4454,6 +4715,15 @@ version = "3.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" +[[package]] +name = "ktx2" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff7f53bdf698e7aa7ec916411bbdc8078135da11b66db5182675b2227f6c0d07" +dependencies = [ + "bitflags 2.9.4", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -4548,6 +4818,15 @@ version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" +[[package]] +name = "mach2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d640282b302c0bb0a2a8e0233ead9035e3bed871f0b7e81fe4a1ec829765db44" +dependencies = [ + "libc", +] + [[package]] name = "malloc_buf" version = "0.0.6" @@ -4626,6 +4905,12 @@ dependencies = [ "paste", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -4661,7 +4946,7 @@ dependencies = [ "log", "num-traits", "pp-rs", - "rustc-hash", + "rustc-hash 1.1.0", "spirv", "termcolor", "thiserror 1.0.69", @@ -4682,7 +4967,7 @@ dependencies = [ "hexf-parse", "indexmap", "log", - "rustc-hash", + "rustc-hash 1.1.0", "spirv", "strum", "termcolor", @@ -4711,7 +4996,7 @@ dependencies = [ "num-traits", "once_cell", "pp-rs", - "rustc-hash", + "rustc-hash 1.1.0", "spirv", "thiserror 2.0.17", "unicode-ident", @@ -4731,7 +5016,7 @@ dependencies = [ "once_cell", "regex", "regex-syntax", - "rustc-hash", + "rustc-hash 1.1.0", "thiserror 1.0.69", "tracing", "unicode-ident", @@ -4748,12 +5033,26 @@ dependencies = [ "indexmap", "naga 26.0.0", "regex", - "rustc-hash", + "rustc-hash 1.1.0", "thiserror 2.0.17", "tracing", "unicode-ident", ] +[[package]] +name = "ndk" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2076a31b7010b17a38c01907c45b945e8f11495ee4dd588309718901b1f7a5b7" +dependencies = [ + "bitflags 2.9.4", + "jni-sys", + "log", + "ndk-sys 0.5.0+25.2.9519653", + "num_enum 0.7.4", + "thiserror 1.0.69", +] + [[package]] name = "ndk" version = "0.9.0" @@ -4822,6 +5121,16 @@ dependencies = [ "rand_xorshift", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nonmax" version = "0.5.5" @@ -4846,6 +5155,17 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "num-derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.106", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -5178,6 +5498,29 @@ dependencies = [ "objc2-foundation 0.2.2", ] +[[package]] +name = "oboe" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8b61bebd49e5d43f5f8cc7ee2891c16e0f41ec7954d36bcb6c14c5e0de867fb" +dependencies = [ + "jni", + "ndk 0.8.0", + "ndk-context", + "num-derive", + "num-traits", + "oboe-sys", +] + +[[package]] +name = "oboe-sys" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c8bb09a4a2b1d668170cfe0a7d5bc103f8999fb316c98099b6a9939c9f2e79d" +dependencies = [ + "cc", +] + [[package]] name = "offset-allocator" version = "0.2.0" @@ -5420,6 +5763,19 @@ dependencies = [ "skeptic", ] +[[package]] +name = "png" +version = "0.17.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526" +dependencies = [ + "bitflags 1.3.2", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + [[package]] name = "png" version = "0.18.0" @@ -5845,6 +6201,15 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "rodio" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7ceb6607dd738c99bc8cb28eff249b7cd5c8ec88b9db96c0608c1480d140fb1" +dependencies = [ + "cpal", +] + [[package]] name = "ron" version = "0.8.1" @@ -5898,6 +6263,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustix" version = "0.38.44" @@ -7074,7 +7445,7 @@ dependencies = [ "parking_lot", "profiling", "raw-window-handle", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 1.0.69", "web-sys", @@ -7100,7 +7471,7 @@ dependencies = [ "parking_lot", "profiling", "raw-window-handle", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 2.0.17", "wgpu-hal 24.0.4", @@ -7128,7 +7499,7 @@ dependencies = [ "portable-atomic", "profiling", "raw-window-handle", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 2.0.17", "wgpu-core-deps-apple", @@ -7211,7 +7582,7 @@ dependencies = [ "range-alloc", "raw-window-handle", "renderdoc-sys", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 1.0.69", "wasm-bindgen", @@ -7256,7 +7627,7 @@ dependencies = [ "range-alloc", "raw-window-handle", "renderdoc-sys", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 2.0.17", "wasm-bindgen", @@ -7893,7 +8264,7 @@ dependencies = [ "dpi", "js-sys", "libc", - "ndk", + "ndk 0.9.0", "objc2 0.5.2", "objc2-app-kit 0.2.2", "objc2-foundation 0.2.2", diff --git a/Cargo.toml b/Cargo.toml index fa493093..2f750188 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -134,6 +134,37 @@ viewer = [ "bevy/png", ] +relighting_example = [ + "bevy/asset_processor", + "bevy/bevy_asset", + "bevy/bevy_audio", + "bevy/bevy_color", + "bevy/bevy_core_pipeline", + "bevy/bevy_post_process", + "bevy/bevy_anti_alias", + "bevy/bevy_gizmos", + "bevy/bevy_gltf", + "bevy/bevy_log", + "bevy/bevy_pbr", + "bevy/bevy_render", + "bevy/bevy_scene", + "bevy/bevy_image", + "bevy/bevy_mesh", + "bevy/bevy_camera", + "bevy/bevy_light", + "bevy/bevy_state", + "bevy/bevy_ui", + "bevy/bevy_text", + "bevy/png", + "bevy/jpeg", + "bevy/compressed_image_saver", + "bevy/reflect_auto_register", + "bevy/reflect_functions", + "bevy/serialize", + "bevy/smaa_luts", + "bevy/tonemapping_luts", +] + web = [ "buffer_storage", "sh0", @@ -167,6 +198,7 @@ bevy_interleave = { version = "0.8.0" } bevy_panorbit_camera = { version = "0.29.0", optional = true , features = ["bevy_egui"]} bevy_transform_gizmo = { version = "0.12", optional = true } bevy_file_asset = { version = "0.3", optional = true } +bevy_pbr = "0.17" bincode2 = { version = "2.0", optional = true } byte-unit = { version = "5.1", optional = true } bytemuck = "1.23" @@ -177,6 +209,7 @@ half = { version = "2.6", features = ["serde"] } # image = { version = "0.25.6", default-features = false, features = ["png"] } kd-tree = { version = "0.6", optional = true } noise = { version = "0.9.0", optional = true } +png = { version = "0.17", optional = true } ply-rs = { version = "0.1", optional = true } rand = "0.8" rayon = { version = "1.8", optional = true } @@ -270,6 +303,11 @@ name = "surfel_plane" path = "tools/surfel_plane.rs" required-features = ["debug_tooling"] +[[bin]] +name = "generate_ggx_lut" +path = "tools/generate_ggx_lut.rs" +required-features = ["png"] + [[bin]] name = "test_gaussian" @@ -293,6 +331,11 @@ path = "examples/headless.rs" name = "multi_camera" path = "examples/multi_camera.rs" +[[example]] +name = "relighting" +path = "examples/relighting.rs" +required-features = ["relighting_example"] + [[bench]] name = "io" diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_AmbientOcclusion.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_AmbientOcclusion.png new file mode 100644 index 00000000..bfddbddf Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_AmbientOcclusion.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_BaseColor.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_BaseColor.png new file mode 100644 index 00000000..f9aeac25 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_BaseColor.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Displacement.tiff b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Displacement.tiff new file mode 100644 index 00000000..544fa1b1 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Displacement.tiff differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Metallic.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Metallic.png new file mode 100644 index 00000000..046578d5 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Metallic.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Normal.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Normal.png new file mode 100644 index 00000000..cbcd5ea1 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Normal.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_ORM.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_ORM.png new file mode 100644 index 00000000..7f5eaa53 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_ORM.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Preview1.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Preview1.png new file mode 100644 index 00000000..08031eef Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Preview1.png differ diff --git a/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Roughness.png b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Roughness.png new file mode 100644 index 00000000..10ea2de6 Binary files /dev/null and b/assets/pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_Roughness.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_AmbientOcclusion.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_AmbientOcclusion.png new file mode 100644 index 00000000..c1701c65 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_AmbientOcclusion.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_BaseColor.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_BaseColor.png new file mode 100644 index 00000000..90a7bf73 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_BaseColor.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Displacement.tiff b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Displacement.tiff new file mode 100644 index 00000000..a933b5b3 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Displacement.tiff differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Metallic.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Metallic.png new file mode 100644 index 00000000..318dca63 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Metallic.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Normal.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Normal.png new file mode 100644 index 00000000..c5f03d2a Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Normal.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_ORM.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_ORM.png new file mode 100644 index 00000000..c94bd6d7 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_ORM.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Preview1.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Preview1.png new file mode 100644 index 00000000..f25a5f4c Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Preview1.png differ diff --git a/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Roughness.png b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Roughness.png new file mode 100644 index 00000000..29038b73 Binary files /dev/null and b/assets/pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_Roughness.png differ diff --git a/examples/relighting.rs b/examples/relighting.rs new file mode 100644 index 00000000..a01be86e --- /dev/null +++ b/examples/relighting.rs @@ -0,0 +1,402 @@ +//! Demonstrates realtime dynamic raytraced lighting with Gaussian Splatting. + +#[path = "relighting/camera_controller.rs"] +mod camera_controller; + +use bevy::{ + camera::CameraMainTextureUsages, gltf::GltfMaterialName, log::LogPlugin, prelude::*, + render::render_resource::TextureUsages, scene::SceneInstanceReady +}; +use bevy_gaussian_splatting::{ + CloudSettings, GaussianCamera, GaussianSplattingPlugin, PlanarGaussian3dHandle, RasterizeMode, +}; +use bevy_gaussian_splatting::material::gaussian_material::{ + GaussianMaterial, GaussianMaterialHandle, GaussianTextureProjection, +}; +use camera_controller::{CameraController, CameraControllerPlugin}; +use std::f32::consts::PI; + +#[derive(Component)] +struct DioramaTag; + +#[derive(Resource, Clone, Copy)] +struct IcecreamPos(pub Vec3); + +#[derive(Resource)] +struct GaussianMaterialCycle { + material_handle: Handle, + textures: Vec>, + current: usize, + timer: Timer, +} + +const TEXTURE_SWAP_SECONDS: f32 = 10.0; + +fn main() { + let mut app = App::new(); + + app.add_plugins(( + DefaultPlugins.set(LogPlugin{ + filter: "wgpu=error,naga=warn,bevy_gaussian_splatting=debug,bevy_render=info,bevy_asset=info".to_string(), + ..default() + }), + GaussianSplattingPlugin, + CameraControllerPlugin + )) + .add_systems(Startup, setup) + .add_systems( + Update, + ( + pause_scene, + toggle_lights, + patrol_path, + cycle_gaussian_material_textures, + ), + ) + .add_systems(PostUpdate, update_text) + .run(); +} + +fn setup( + mut commands: Commands, + asset_server: Res, + mut meshes: ResMut>, + mut materials: ResMut>, + mut gaussian_materials: ResMut>, +) { + let diorama_scene: Handle = asset_server.load( + GltfAssetLabel::Scene(0).from_asset( + "https://github.com/bevyengine/bevy_asset_files/raw/2a5950295a8b6d9d051d59c0df69e87abcda58c3/pica_pica/mini_diorama_01.glb", + ), + ); + commands + .spawn((SceneRoot(diorama_scene), Transform::from_scale(Vec3::splat(10.0)), DioramaTag)) + .observe(add_mesh_processing_on_scene_load); + + commands + .spawn(( + SceneRoot(asset_server.load( + GltfAssetLabel::Scene(0).from_asset("https://github.com/bevyengine/bevy_asset_files/raw/2a5950295a8b6d9d051d59c0df69e87abcda58c3/pica_pica/robot_01.glb") + )), + Transform::from_scale(Vec3::splat(2.0)) + .with_translation(Vec3::new(-2.0, 0.05, -2.1)) + .with_rotation(Quat::from_rotation_y(PI / 2.0)), + PatrolPath { + path: vec![ + (Vec3::new(-2.0, 0.05, -2.1), Quat::from_rotation_y(PI / 2.0)), + (Vec3::new(2.2, 0.05, -2.1), Quat::from_rotation_y(0.0)), + ( + Vec3::new(2.2, 0.05, 2.1), + Quat::from_rotation_y(3.0 * PI / 2.0), + ), + (Vec3::new(-2.0, 0.05, 2.1), Quat::from_rotation_y(PI)), + ], + i: 0, + }, + )) + .observe(add_mesh_processing_on_scene_load); + + let icecream_pos = Vec3::new(0.0, 1.0, 0.0); + + let gaussian_textures = vec![ + asset_server.load("pbr/Poliigon_BrickWallReclaimed_8320/Poliigon_BrickWallReclaimed_8320_BaseColor.png"), + asset_server.load("pbr/Poliigon_GrassPatchyGround_4585/Poliigon_GrassPatchyGround_4585_BaseColor.png"), + ]; + + let gaussian_material_handle = gaussian_materials.add(GaussianMaterial { + base_color: LinearRgba::WHITE, + base_color_texture: gaussian_textures.first().cloned(), + texture_projection: GaussianTextureProjection::Xz, + bounds: None, + }); + + commands.spawn(( + PlanarGaussian3dHandle( + asset_server.load( + "https://raw.githubusercontent.com/mosure/bevy_gaussian_splatting/main/assets/scenes/icecream.ply", + ), + ), + CloudSettings { + aabb: true, + visualize_bounding_box: false, + global_opacity: 6.0, + rasterize_mode: RasterizeMode::Color, + sort_mode: bevy_gaussian_splatting::sort::SortMode::Std, + ..default() + }, + Transform::from_translation(icecream_pos).with_scale(Vec3::splat(1.5)), + GaussianMaterialHandle(gaussian_material_handle.clone()), + )); + + commands.insert_resource(IcecreamPos(icecream_pos)); + commands.insert_resource(GaussianMaterialCycle { + material_handle: gaussian_material_handle, + textures: gaussian_textures, + current: 0, + timer: Timer::from_seconds(TEXTURE_SWAP_SECONDS, TimerMode::Repeating), + }); + let marker_mesh = meshes.add(Cuboid::new(0.2, 0.2, 0.2)); + let marker_mat = materials.add(StandardMaterial { + base_color: Color::srgb(1.0, 0.2, 0.2), + emissive: LinearRgba::from(Color::srgb(1.0, 0.1, 0.1)) * 50.0, + ..default() + }); + commands.spawn((Mesh3d(marker_mesh), MeshMaterial3d(marker_mat), Transform::from_translation(icecream_pos))); + + commands.spawn(( + DirectionalLight { + illuminance: light_consts::lux::FULL_DAYLIGHT, + shadows_enabled: false, + ..default() + }, + Transform::from_rotation(Quat::from_xyzw( + -0.13334629, + -0.86597735, + -0.3586996, + 0.3219264, + )), + )); + + let eye = Vec3::new(0.0, 2.0, 5.0); + commands.spawn(( + GaussianCamera { warmup: false }, + Camera3d::default(), + Camera { + clear_color: ClearColorConfig::Custom(Color::BLACK), + ..default() + }, + CameraController { + walk_speed: 3.0, + run_speed: 10.0, + ..Default::default() + }, + Transform::from_translation(eye).looking_at(icecream_pos, Vec3::Y), + CameraMainTextureUsages::default().with(TextureUsages::STORAGE_BINDING), + Msaa::Off, + )); + + commands.spawn(( + Text::new("Loading..."), + TextFont { + font_size: 16.0, + ..default() + }, + TextColor(Color::WHITE), + Node { + position_type: PositionType::Absolute, + bottom: Val::Px(12.0), + left: Val::Px(12.0), + ..default() + }, + )); +} + +fn add_mesh_processing_on_scene_load( + scene_ready: On, + children: Query<&Children>, + mesh_query: Query<( + &Mesh3d, + &MeshMaterial3d, + Option<&GltfMaterialName>, + )>, + mut meshes: ResMut>, + mut materials: ResMut>, + mut commands: Commands, +) { + for descendant in children.iter_descendants(scene_ready.entity) { + if let Ok((Mesh3d(mesh_handle), MeshMaterial3d(material_handle), material_name)) = + mesh_query.get(descendant) + { + let mesh = meshes.get_mut(mesh_handle).unwrap(); + if !mesh.contains_attribute(Mesh::ATTRIBUTE_UV_0) { + let vertex_count = mesh.count_vertices(); + mesh.insert_attribute(Mesh::ATTRIBUTE_UV_0, vec![[0.0, 0.0]; vertex_count]); + mesh.insert_attribute( + Mesh::ATTRIBUTE_TANGENT, + vec![[0.0, 0.0, 0.0, 0.0]; vertex_count], + ); + } + if !mesh.contains_attribute(Mesh::ATTRIBUTE_TANGENT) { + mesh.generate_tangents().unwrap(); + } + if mesh.contains_attribute(Mesh::ATTRIBUTE_UV_1) { + mesh.remove_attribute(Mesh::ATTRIBUTE_UV_1); + } + + if material_name.map(|s| s.0.as_str()) == Some("material") { + let material = materials.get_mut(material_handle).unwrap(); + material.emissive = LinearRgba::BLACK; + } + if material_name.map(|s| s.0.as_str()) == Some("Lights") { + let material = materials.get_mut(material_handle).unwrap(); + material.emissive = + LinearRgba::from(Color::srgb(0.941, 0.714, 0.043)) * 1_000_000.0; + material.alpha_mode = AlphaMode::Opaque; + material.specular_transmission = 0.0; + + commands.insert_resource(RobotLightMaterial(material_handle.clone())); + } + if material_name.map(|s| s.0.as_str()) == Some("Glass_Dark_01") { + let material = materials.get_mut(material_handle).unwrap(); + material.alpha_mode = AlphaMode::Opaque; + material.specular_transmission = 0.0; + } + } + } +} + +fn pause_scene(mut time: ResMut>, key_input: Res>) { + if key_input.just_pressed(KeyCode::Space) { + if time.is_paused() { + time.unpause(); + } else { + time.pause(); + } + } +} + +#[derive(Resource)] +struct RobotLightMaterial(Handle); + +fn toggle_lights( + key_input: Res>, + robot_light_material: Option>, + mut materials: ResMut>, + directional_light: Query>, + diorama: Query>, + icecream_pos: Option>, + mut cam_q: Query<&mut Transform, With>, + mut commands: Commands, +) { + if key_input.just_pressed(KeyCode::Digit1) { + if let Ok(directional_light) = directional_light.single() { + commands.entity(directional_light).despawn(); + } else { + commands.spawn(( + DirectionalLight { + illuminance: light_consts::lux::FULL_DAYLIGHT, + shadows_enabled: false, + ..default() + }, + Transform::from_rotation(Quat::from_xyzw( + -0.13334629, + -0.86597735, + -0.3586996, + 0.3219264, + )), + )); + } + } + + if key_input.just_pressed(KeyCode::Digit3) { + if let Ok(entity) = diorama.single() { + commands.entity(entity).despawn(); + } + } + + if key_input.just_pressed(KeyCode::Digit4) { + if let (Some(pos), Ok(mut cam_tf)) = (icecream_pos, cam_q.single_mut()) { + let target = pos.0; + let eye = target + Vec3::new(0.0, 0.6, 2.6); + *cam_tf = Transform::from_translation(eye).looking_at(target, Vec3::Y); + } + } + + if key_input.just_pressed(KeyCode::Digit2) + && let Some(robot_light_material) = robot_light_material + { + let material = materials.get_mut(&robot_light_material.0).unwrap(); + if material.emissive == LinearRgba::BLACK { + material.emissive = LinearRgba::from(Color::srgb(0.941, 0.714, 0.043)) * 1_000_000.0; + } else { + material.emissive = LinearRgba::BLACK; + } + } +} + +#[derive(Component)] +struct PatrolPath { + path: Vec<(Vec3, Quat)>, + i: usize, +} + +fn patrol_path(mut query: Query<(&mut PatrolPath, &mut Transform)>, time: Res>) { + for (mut path, mut transform) in query.iter_mut() { + let (mut target_position, mut target_rotation) = path.path[path.i]; + let mut distance_to_target = transform.translation.distance(target_position); + if distance_to_target < 0.01 { + transform.translation = target_position; + transform.rotation = target_rotation; + + path.i = (path.i + 1) % path.path.len(); + (target_position, target_rotation) = path.path[path.i]; + distance_to_target = transform.translation.distance(target_position); + } + + let direction = (target_position - transform.translation).normalize(); + let movement = direction * time.delta_secs(); + + if movement.length() > distance_to_target { + transform.translation = target_position; + transform.rotation = target_rotation; + } else { + transform.translation += movement; + } + } +} + +fn update_text( + mut text: Single<&mut Text>, + robot_light_material: Option>, + materials: Res>, + directional_light: Query>, + time: Res>, +) { + let mut content = String::new(); + + if time.is_paused() { + content.push_str("(Space): Resume"); + } else { + content.push_str("(Space): Pause"); + } + + if directional_light.single().is_ok() { + content.push_str("\n(1): Disable directional light"); + } else { + content.push_str("\n(1): Enable directional light"); + } + + content.push_str("\n(3): Despawn diorama (unhide icecream)"); + content.push_str("\n(4): Teleport camera to icecream"); + + match robot_light_material.and_then(|m| materials.get(&m.0)) { + Some(robot_light_material) if robot_light_material.emissive != LinearRgba::BLACK => { + content.push_str("\n(2): Disable robot emissive light"); + } + _ => { + content.push_str("\n(2): Enable robot emissive light"); + } + } + + text.0 = content; +} + +fn cycle_gaussian_material_textures( + time: Res>, + mut gaussian_materials: ResMut>, + mut cycle: ResMut, +) { + if cycle.textures.is_empty() { + return; + } + + if !cycle.timer.tick(time.delta()).just_finished() { + return; + } + + cycle.current = (cycle.current + 1) % cycle.textures.len(); + + if let Some(material) = gaussian_materials.get_mut(&cycle.material_handle) { + material.base_color_texture = Some(cycle.textures[cycle.current].clone()); + } +} diff --git a/examples/relighting/camera_controller.rs b/examples/relighting/camera_controller.rs new file mode 100644 index 00000000..7c9b4191 --- /dev/null +++ b/examples/relighting/camera_controller.rs @@ -0,0 +1,218 @@ +//! A freecam-style camera controller plugin. + +use bevy::{ + input::mouse::{AccumulatedMouseMotion, AccumulatedMouseScroll, MouseScrollUnit}, + prelude::*, + window::{CursorGrabMode, CursorOptions}, +}; +use std::{f32::consts::*, fmt}; + +pub struct CameraControllerPlugin; + +impl Plugin for CameraControllerPlugin { + fn build(&self, app: &mut App) { + app.add_systems(Update, run_camera_controller); + } +} + +pub const RADIANS_PER_DOT: f32 = 1.0 / 180.0; + +#[derive(Component)] +pub struct CameraController { + pub enabled: bool, + pub initialized: bool, + pub sensitivity: f32, + pub key_forward: KeyCode, + pub key_back: KeyCode, + pub key_left: KeyCode, + pub key_right: KeyCode, + pub key_up: KeyCode, + pub key_down: KeyCode, + pub key_run: KeyCode, + pub mouse_key_cursor_grab: MouseButton, + pub keyboard_key_toggle_cursor_grab: KeyCode, + pub walk_speed: f32, + pub run_speed: f32, + pub scroll_factor: f32, + pub friction: f32, + pub pitch: f32, + pub yaw: f32, + pub velocity: Vec3, +} + +impl Default for CameraController { + fn default() -> Self { + Self { + enabled: true, + initialized: false, + sensitivity: 1.0, + key_forward: KeyCode::KeyW, + key_back: KeyCode::KeyS, + key_left: KeyCode::KeyA, + key_right: KeyCode::KeyD, + key_up: KeyCode::KeyE, + key_down: KeyCode::KeyQ, + key_run: KeyCode::ShiftLeft, + mouse_key_cursor_grab: MouseButton::Left, + keyboard_key_toggle_cursor_grab: KeyCode::KeyM, + walk_speed: 5.0, + run_speed: 15.0, + scroll_factor: 0.1, + friction: 0.5, + pitch: 0.0, + yaw: 0.0, + velocity: Vec3::ZERO, + } + } +} + +impl fmt::Display for CameraController { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + " +Freecam Controls: + Mouse\t- Move camera orientation + Scroll\t- Adjust movement speed + {:?}\t- Hold to grab cursor + {:?}\t- Toggle cursor grab + {:?} & {:?}\t- Fly forward & backwards + {:?} & {:?}\t- Fly sideways left & right + {:?} & {:?}\t- Fly up & down + {:?}\t- Fly faster while held", + self.mouse_key_cursor_grab, + self.keyboard_key_toggle_cursor_grab, + self.key_forward, + self.key_back, + self.key_left, + self.key_right, + self.key_up, + self.key_down, + self.key_run, + ) + } +} + +fn run_camera_controller( + time: Res>, + mut windows: Query<(&Window, &mut CursorOptions)>, + accumulated_mouse_motion: Res, + accumulated_mouse_scroll: Res, + mouse_button_input: Res>, + key_input: Res>, + mut toggle_cursor_grab: Local, + mut mouse_cursor_grab: Local, + mut query: Query<(&mut Transform, &mut CameraController), With>, +) { + let dt = time.delta_secs(); + + let Ok((mut transform, mut controller)) = query.single_mut() else { + return; + }; + + if !controller.initialized { + let (yaw, pitch, _roll) = transform.rotation.to_euler(EulerRot::YXZ); + controller.yaw = yaw; + controller.pitch = pitch; + controller.initialized = true; + info!("{}", *controller); + } + if !controller.enabled { + return; + } + + let mut scroll = 0.0; + + let amount = match accumulated_mouse_scroll.unit { + MouseScrollUnit::Line => accumulated_mouse_scroll.delta.y, + MouseScrollUnit::Pixel => accumulated_mouse_scroll.delta.y / 16.0, + }; + scroll += amount; + controller.walk_speed += scroll * controller.scroll_factor * controller.walk_speed; + controller.run_speed = controller.walk_speed * 3.0; + + let mut axis_input = Vec3::ZERO; + if key_input.pressed(controller.key_forward) { + axis_input.z += 1.0; + } + if key_input.pressed(controller.key_back) { + axis_input.z -= 1.0; + } + if key_input.pressed(controller.key_right) { + axis_input.x += 1.0; + } + if key_input.pressed(controller.key_left) { + axis_input.x -= 1.0; + } + if key_input.pressed(controller.key_up) { + axis_input.y += 1.0; + } + if key_input.pressed(controller.key_down) { + axis_input.y -= 1.0; + } + + let mut cursor_grab_change = false; + if key_input.just_pressed(controller.keyboard_key_toggle_cursor_grab) { + *toggle_cursor_grab = !*toggle_cursor_grab; + cursor_grab_change = true; + } + if mouse_button_input.just_pressed(controller.mouse_key_cursor_grab) { + *mouse_cursor_grab = true; + cursor_grab_change = true; + } + if mouse_button_input.just_released(controller.mouse_key_cursor_grab) { + *mouse_cursor_grab = false; + cursor_grab_change = true; + } + let cursor_grab = *mouse_cursor_grab || *toggle_cursor_grab; + + if axis_input != Vec3::ZERO { + let max_speed = if key_input.pressed(controller.key_run) { + controller.run_speed + } else { + controller.walk_speed + }; + controller.velocity = axis_input.normalize() * max_speed; + } else { + let friction = controller.friction.clamp(0.0, 1.0); + controller.velocity *= 1.0 - friction; + if controller.velocity.length_squared() < 1e-6 { + controller.velocity = Vec3::ZERO; + } + } + + if controller.velocity != Vec3::ZERO { + let forward = *transform.forward(); + let right = *transform.right(); + transform.translation += controller.velocity.x * dt * right + + controller.velocity.y * dt * Vec3::Y + + controller.velocity.z * dt * forward; + } + + if cursor_grab_change { + if cursor_grab { + for (window, mut cursor_options) in &mut windows { + if !window.focused { + continue; + } + + cursor_options.grab_mode = CursorGrabMode::Locked; + cursor_options.visible = false; + } + } else { + for (_, mut cursor_options) in &mut windows { + cursor_options.grab_mode = CursorGrabMode::None; + cursor_options.visible = true; + } + } + } + + if accumulated_mouse_motion.delta != Vec2::ZERO && cursor_grab { + controller.pitch = (controller.pitch + - accumulated_mouse_motion.delta.y * RADIANS_PER_DOT * controller.sensitivity) + .clamp(-PI / 2., PI / 2.); + controller.yaw -= + accumulated_mouse_motion.delta.x * RADIANS_PER_DOT * controller.sensitivity; + transform.rotation = Quat::from_euler(EulerRot::ZYX, 0.0, controller.yaw, controller.pitch); + } +} diff --git a/src/lib.rs b/src/lib.rs index 0a1d716c..b075db42 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,7 @@ pub mod io; pub mod material; pub mod math; pub mod morph; +pub mod pbr_decomposition; pub mod query; pub mod render; pub mod sort; @@ -62,7 +63,11 @@ impl Plugin for GaussianSplattingPlugin { render::RenderPipelinePlugin::::default(), )); - app.add_plugins((material::MaterialPlugin, query::QueryPlugin)); + app.add_plugins(( + material::MaterialPlugin, + pbr_decomposition::PbrDecompositionPlugin, + query::QueryPlugin, + )); #[cfg(feature = "noise")] app.add_plugins(noise::NoisePlugin); diff --git a/src/material/gaussian_material.rs b/src/material/gaussian_material.rs new file mode 100644 index 00000000..a7aff326 --- /dev/null +++ b/src/material/gaussian_material.rs @@ -0,0 +1,214 @@ +use std::collections::{HashMap, HashSet}; + +use bevy::{ + asset::Asset, + ecs::query::QueryItem, + prelude::*, + render::{ + extract_component::{ExtractComponent, ExtractComponentPlugin}, + render_asset::RenderAssets, + render_resource::{Sampler, SamplerId, TextureView, TextureViewId}, + Extract, ExtractSchedule, Render, RenderApp, RenderSystems, + }, +}; +use bevy::render::texture::{FallbackImage, GpuImage}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)] +#[reflect(Default)] +pub enum GaussianTextureProjection { + Xy, + Xz, + Yz, +} + +impl Default for GaussianTextureProjection { + fn default() -> Self { + Self::Xz + } +} + +#[derive(Default, Clone, Copy, Debug, Serialize, Deserialize, Reflect, PartialEq)] +#[reflect(Default)] +pub struct GaussianMaterialBounds { + pub min: Vec3, + pub max: Vec3, +} + +#[derive(Asset, Clone, Debug, Reflect)] +#[reflect(Default)] +pub struct GaussianMaterial { + pub base_color: LinearRgba, + pub base_color_texture: Option>, + pub texture_projection: GaussianTextureProjection, + /// Optional manual bounds. When `None`, the gaussian cloud bounds will be used. + pub bounds: Option, +} + +impl Default for GaussianMaterial { + fn default() -> Self { + Self { + base_color: LinearRgba::WHITE, + base_color_texture: None, + texture_projection: GaussianTextureProjection::default(), + bounds: None, + } + } +} + +#[derive(Clone, Debug)] +pub struct GpuGaussianMaterial { + pub base_color: Vec4, + pub projection_axis: u32, + pub bounds_override: Option, + pub texture_view: TextureView, + pub sampler: Sampler, + pub use_texture: u32, + pub texture_view_id: TextureViewId, + pub sampler_id: SamplerId, +} + +impl PartialEq for GpuGaussianMaterial { + fn eq(&self, other: &Self) -> bool { + self.base_color == other.base_color + && self.projection_axis == other.projection_axis + && self.bounds_override == other.bounds_override + && self.use_texture == other.use_texture + && self.texture_view_id == other.texture_view_id + && self.sampler_id == other.sampler_id + } +} + +#[derive(Clone, Debug)] +pub struct CachedGaussianMaterial { + pub material: GpuGaussianMaterial, + pub revision: u64, +} + +#[derive(Resource, Default)] +pub struct RenderGaussianMaterials { + pub map: HashMap, CachedGaussianMaterial>, + pub revision: u64, +} + +#[derive(Resource, Default, Clone)] +pub struct ExtractedGaussianMaterials(pub Vec<(AssetId, GaussianMaterial)>); + +#[derive(Component, Clone, Reflect, Deref, DerefMut)] +#[reflect(Component, Default)] +pub struct GaussianMaterialHandle(pub Handle); + +impl Default for GaussianMaterialHandle { + fn default() -> Self { + Self(Handle::default()) + } +} + +impl ExtractComponent for GaussianMaterialHandle { + type QueryData = &'static Self; + type QueryFilter = (); + type Out = Self; + + fn extract_component(component: QueryItem<'_, '_, Self::QueryData>) -> Option { + Some(component.clone()) + } +} + +pub struct GaussianMaterialPlugin; + +impl Plugin for GaussianMaterialPlugin { + fn build(&self, app: &mut App) { + app.register_type::() + .register_type::() + .register_type::(); + app.init_asset::(); + app.add_plugins(ExtractComponentPlugin::::default()); + + if let Some(render_app) = app.get_sub_app_mut(RenderApp) { + render_app + .init_resource::() + .init_resource::() + .add_systems(ExtractSchedule, extract_gaussian_materials) + .add_systems( + Render, + prepare_gaussian_materials_system.in_set(RenderSystems::PrepareBindGroups), + ); + } + } +} + +fn prepare_gaussian_materials_system( + materials: Res, + images: Res>, + fallback_image: Res, + mut cache: ResMut, +) { + let fallback_texture_view = fallback_image.d2.texture_view.clone(); + let fallback_sampler = fallback_image.d2.sampler.clone(); + let mut seen_ids: HashSet> = + HashSet::with_capacity(materials.0.len()); + let mut revision = cache.revision; + + for (id, material) in materials.0.iter() { + seen_ids.insert(*id); + let (texture_view, sampler, use_texture) = if let Some(handle) = material.base_color_texture.clone() { + if let Some(image) = images.get(&handle) { + (image.texture_view.clone(), image.sampler.clone(), 1) + } else { + (fallback_texture_view.clone(), fallback_sampler.clone(), 0) + } + } else { + (fallback_texture_view.clone(), fallback_sampler.clone(), 0) + }; + + let projection_axis = match material.texture_projection { + GaussianTextureProjection::Xy => 0, + GaussianTextureProjection::Xz => 1, + GaussianTextureProjection::Yz => 2, + }; + + let texture_view_id = texture_view.id(); + let sampler_id = sampler.id(); + + let gpu_material = GpuGaussianMaterial { + base_color: material.base_color.to_vec4(), + projection_axis, + bounds_override: material.bounds, + texture_view, + sampler, + use_texture, + texture_view_id, + sampler_id, + }; + + let entry = cache.map.get(id).cloned(); + + match entry { + Some(existing) if existing.material == gpu_material => {} + _ => { + revision = revision.wrapping_add(1); + cache.map.insert( + *id, + CachedGaussianMaterial { + material: gpu_material, + revision, + }, + ); + } + } + } + + cache.map.retain(|id, _| seen_ids.contains(id)); + cache.revision = revision; +} + +fn extract_gaussian_materials( + mut commands: Commands, + materials: Extract>>, +) { + let mut extracted = Vec::with_capacity(materials.len()); + for (id, material) in materials.iter() { + extracted.push((id, material.clone())); + } + commands.insert_resource(ExtractedGaussianMaterials(extracted)); +} diff --git a/src/material/mod.rs b/src/material/mod.rs index 73e6b60e..e6f566a5 100644 --- a/src/material/mod.rs +++ b/src/material/mod.rs @@ -6,6 +6,7 @@ pub mod optical_flow; pub mod position; pub mod spherical_harmonics; pub mod spherindrical_harmonics; +pub mod gaussian_material; #[cfg(feature = "material_noise")] pub mod noise; @@ -26,6 +27,7 @@ impl Plugin for MaterialPlugin { position::PositionMaterialPlugin, spherical_harmonics::SphericalHarmonicCoefficientsPlugin, spherindrical_harmonics::SpherindricalHarmonicCoefficientsPlugin, + gaussian_material::GaussianMaterialPlugin, )); } } diff --git a/src/pbr_decomposition/energy_calibration.rs b/src/pbr_decomposition/energy_calibration.rs new file mode 100644 index 00000000..26a6fc8e --- /dev/null +++ b/src/pbr_decomposition/energy_calibration.rs @@ -0,0 +1 @@ +use bevy::prelude::*; diff --git a/src/pbr_decomposition/material_separation.rs b/src/pbr_decomposition/material_separation.rs new file mode 100644 index 00000000..70a559f5 --- /dev/null +++ b/src/pbr_decomposition/material_separation.rs @@ -0,0 +1,5 @@ +use bevy::prelude::*; + +pub mod pipeline; + +pub use pipeline::*; diff --git a/src/pbr_decomposition/material_separation/pipeline.rs b/src/pbr_decomposition/material_separation/pipeline.rs new file mode 100644 index 00000000..4ecd5c88 --- /dev/null +++ b/src/pbr_decomposition/material_separation/pipeline.rs @@ -0,0 +1,138 @@ +use bevy::{ + asset::{load_internal_asset, uuid_handle}, + prelude::*, + render::{ + render_resource::*, + renderer::RenderDevice, + }, +}; +use std::collections::HashMap; +use bevy::asset::AssetId; +use crate::gaussian::formats::planar_3d::PlanarGaussian3d; + +use crate::pbr_decomposition::types::*; +use crate::render::{CloudPipeline, CloudPipelineKey, shader_defs}; +use crate::gaussian::formats::planar_3d::Gaussian3d; + +const MATERIAL_SEPARATION_SHADER_HANDLE: Handle = + uuid_handle!("d4e5f6a7-b8c9-0123-def0-234567890123"); + +#[derive(Resource)] +pub struct MaterialSeparationBuffers { + pub asset_map: HashMap, GpuMaterialBuffers>, +} + +impl Default for MaterialSeparationBuffers { + fn default() -> Self { + MaterialSeparationBuffers { + asset_map: HashMap::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct GpuMaterialBuffers { + pub materials: Buffer, +} + +impl GpuMaterialBuffers { + pub fn new(count: usize, render_device: &RenderDevice) -> Self { + let materials = render_device.create_buffer(&BufferDescriptor { + label: Some("pbr_materials"), + size: (count as u64) * PbrMaterialData::min_size().get(), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + Self { materials } + } +} + +#[derive(Resource)] +pub struct MaterialSeparationPipeline { + pub pipeline: CachedComputePipelineId, + pub bind_group_layout: BindGroupLayout, + pub settings_layout: BindGroupLayout, +} + +impl FromWorld for MaterialSeparationPipeline { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + let bind_group_layout = render_device.create_bind_group_layout( + Some("material_separation_bind_group_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: Some(StreamingStats::min_size()), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: Some(PbrMaterialData::min_size()), + }, + count: None, + }, + ], + ); + + let settings_layout = render_device.create_bind_group_layout( + Some("material_separation_settings_layout"), + &[BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }], + ); + + let pipeline_cache = world.resource::(); + let cloud_pipeline = world.resource::>(); + let shader_defs = shader_defs(CloudPipelineKey::default()); + + let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("material_separation".into()), + layout: vec![ + cloud_pipeline.compute_view_layout.clone(), + cloud_pipeline.gaussian_uniform_layout.clone(), + cloud_pipeline.gaussian_cloud_layout.clone(), + bind_group_layout.clone(), + settings_layout.clone(), + cloud_pipeline.gaussian_material_layout.clone(), + ], + push_constant_ranges: vec![], + shader: MATERIAL_SEPARATION_SHADER_HANDLE, + shader_defs, + entry_point: Some("estimate_material_properties".into()), + zero_initialize_workgroup_memory: false, + }); + + Self { + pipeline, + bind_group_layout, + settings_layout, + } + } +} + +pub fn load_material_separation_shader(app: &mut App) { + load_internal_asset!( + app, + MATERIAL_SEPARATION_SHADER_HANDLE, + "../shaders/material_separation.wgsl", + Shader::from_wgsl + ); +} diff --git a/src/pbr_decomposition/mod.rs b/src/pbr_decomposition/mod.rs new file mode 100644 index 00000000..b8f5ec39 --- /dev/null +++ b/src/pbr_decomposition/mod.rs @@ -0,0 +1,78 @@ +use bevy::prelude::*; +use bevy::render::{RenderApp, Render}; +use bevy::core_pipeline::core_3d::graph::{Core3d, Node3d}; +use bevy::render::renderer::RenderDevice; +use bevy::prelude::Res; +use bevy::prelude::Commands; + +pub mod settings; +pub mod types; +pub mod spatial_hash; +pub mod normal_estimation; +pub mod synthetic_views; +pub mod material_separation; +pub mod energy_calibration; +pub mod orchestrator; + +pub use settings::{PbrDecompositionSettings, SHCoordinateFrame}; +pub use types::*; +pub use orchestrator::*; + +use spatial_hash::load_spatial_hash_shader; +use normal_estimation::load_normal_estimation_shader; +use synthetic_views::load_synthetic_views_shader; +use material_separation::load_material_separation_shader; + +pub struct PbrDecompositionPlugin; + +impl Plugin for PbrDecompositionPlugin { + fn build(&self, app: &mut App) { + app.init_resource::(); + + app.register_type::(); + app.register_type::(); + app.register_type::(); + app.register_type::(); + app.register_type::(); + app.register_type::(); + + load_spatial_hash_shader(app); + load_normal_estimation_shader(app); + load_synthetic_views_shader(app); + load_material_separation_shader(app); + + if let Some(render_app) = app.get_sub_app_mut(RenderApp) { + use orchestrator::{PbrDecompositionLabel, PbrDecompositionNode3d, update_pbr_buffers, extract_settings_to_render}; + use bevy::render::render_graph::RenderGraphExt; + use crate::sort::radix::RadixSortLabel; + + // Mirror settings into render world each frame + render_app.add_systems(bevy::render::ExtractSchedule, extract_settings_to_render); + + // GPU buffer alloc + render_app + .init_resource::() + .init_resource::() + .init_resource::() + .add_systems(Render, (ensure_pipelines, update_pbr_buffers)); + + // Render graph node + render_app.add_render_graph_node::(Core3d, PbrDecompositionLabel); + render_app.add_render_graph_edge(Core3d, RadixSortLabel, PbrDecompositionLabel); + render_app.add_render_graph_edge(Core3d, PbrDecompositionLabel, Node3d::LatePrepass); + } + } +} + +fn ensure_pipelines( + mut commands: Commands, + device: Option>, + has_normals: Option>, + has_synth: Option>, + has_mats: Option>, +) { + if device.is_none() { return; } + if has_normals.is_none() { commands.init_resource::(); } + if has_synth.is_none() { commands.init_resource::(); } + if has_mats.is_none() { commands.init_resource::(); } +} diff --git a/src/pbr_decomposition/normal_estimation.rs b/src/pbr_decomposition/normal_estimation.rs new file mode 100644 index 00000000..70a559f5 --- /dev/null +++ b/src/pbr_decomposition/normal_estimation.rs @@ -0,0 +1,5 @@ +use bevy::prelude::*; + +pub mod pipeline; + +pub use pipeline::*; diff --git a/src/pbr_decomposition/normal_estimation/pipeline.rs b/src/pbr_decomposition/normal_estimation/pipeline.rs new file mode 100644 index 00000000..e7694979 --- /dev/null +++ b/src/pbr_decomposition/normal_estimation/pipeline.rs @@ -0,0 +1,126 @@ +use bevy::{ + asset::{load_internal_asset, uuid_handle}, + prelude::*, + render::{ + render_resource::*, + renderer::RenderDevice, + }, +}; +use std::collections::HashMap; +use bevy::asset::AssetId; +use crate::gaussian::formats::planar_3d::PlanarGaussian3d; + +use crate::pbr_decomposition::types::*; +use crate::render::{CloudPipeline, CloudPipelineKey, shader_defs}; +use crate::gaussian::formats::planar_3d::Gaussian3d; + +const NORMAL_ESTIMATION_SHADER_HANDLE: Handle = + uuid_handle!("b2c3d4e5-f6a7-8901-bcde-f12345678901"); + +#[derive(Resource)] +pub struct NormalEstimationBuffers { + pub asset_map: HashMap, GpuNormalBuffers>, +} + +impl Default for NormalEstimationBuffers { + fn default() -> Self { + NormalEstimationBuffers { + asset_map: HashMap::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct GpuNormalBuffers { + pub normals: Buffer, +} + +impl GpuNormalBuffers { + pub fn new(count: usize, render_device: &RenderDevice) -> Self { + let normals = render_device.create_buffer(&BufferDescriptor { + label: Some("gaussian normals"), + size: (count as u64) * NormalData::min_size().get(), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + Self { normals } + } +} + +#[derive(Resource)] +pub struct NormalEstimationPipeline { + pub pipeline: CachedComputePipelineId, + pub bind_group_layout: BindGroupLayout, + pub settings_layout: BindGroupLayout, +} + +impl FromWorld for NormalEstimationPipeline { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + // Group 3: output normals + let bind_group_layout = render_device.create_bind_group_layout( + Some("normal_estimation_bind_group_layout"), + &[BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: Some(NormalData::min_size()), + }, + count: None, + }], + ); + + let settings_layout = render_device.create_bind_group_layout( + Some("normal_estimation_settings_layout"), + &[BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }], + ); + + let pipeline_cache = world.resource::(); + let cloud_pipeline = world.resource::>(); + let shader_defs = shader_defs(CloudPipelineKey::default()); + + let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("normal_estimation".into()), + layout: vec![ + cloud_pipeline.compute_view_layout.clone(), + cloud_pipeline.gaussian_uniform_layout.clone(), + cloud_pipeline.gaussian_cloud_layout.clone(), + bind_group_layout.clone(), + settings_layout.clone(), + ], + push_constant_ranges: vec![], + shader: NORMAL_ESTIMATION_SHADER_HANDLE, + shader_defs, + entry_point: Some("estimate_normals".into()), + zero_initialize_workgroup_memory: false, + }); + + Self { + pipeline, + bind_group_layout, + settings_layout, + } + } +} + +pub fn load_normal_estimation_shader(app: &mut App) { + load_internal_asset!( + app, + NORMAL_ESTIMATION_SHADER_HANDLE, + "../shaders/normal_estimation.wgsl", + Shader::from_wgsl + ); +} diff --git a/src/pbr_decomposition/orchestrator.rs b/src/pbr_decomposition/orchestrator.rs new file mode 100644 index 00000000..17555b9d --- /dev/null +++ b/src/pbr_decomposition/orchestrator.rs @@ -0,0 +1,315 @@ +use bevy::prelude::*; +use bevy::render::{ + Render, RenderApp, RenderSystems, + render_asset::RenderAssets, + render_graph::{Node, NodeRunError, RenderGraphContext, RenderLabel, RenderGraphExt}, + render_resource::{ + BindGroup, BindGroupEntry, BindingResource, BufferBinding, BufferUsages, BufferInitDescriptor, + ComputePassDescriptor, CachedPipelineState, + }, + renderer::{RenderContext, RenderDevice}, + view::ViewUniformOffset, +}; +use bevy::core_pipeline::prepass::PreviousViewUniformOffset; +use bevy::core_pipeline::core_3d::graph::{Core3d, Node3d}; +use bevy_interleave::{interface::storage::PlanarStorageBindGroup, prelude::{PlanarSync, GpuPlanar}}; +use bytemuck::{Pod, Zeroable}; + +use crate::{ + GaussianCamera, + gaussian::formats::planar_3d::{Gaussian3d, PlanarGaussian3d, PlanarGaussian3dHandle}, + render::{ + CloudPipeline, CloudUniform, GaussianComputeViewBindGroup, GaussianMaterialOverrideBindGroup, + GaussianUniformBindGroups, + }, + sort::radix::RadixSortLabel, +}; +use crate::pbr_decomposition::{ + settings::PbrDecompositionSettings, + normal_estimation::{NormalEstimationPipeline, NormalEstimationBuffers}, + synthetic_views::{SyntheticViewsPipeline, SyntheticViewsBuffers}, + material_separation::{MaterialSeparationPipeline, MaterialSeparationBuffers, GpuMaterialBuffers}, +}; +use crate::pbr_decomposition::synthetic_views::GpuSyntheticViewsBuffers; +use bevy::render::extract_component::DynamicUniformIndex; + +#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Reflect)] +#[reflect(Component)] +pub enum DecompositionStatus { + Pending, + InProgress, + Complete, + Failed, +} + +impl Default for DecompositionStatus { + fn default() -> Self { + Self::Pending + } +} + +#[derive(Component, Debug, Default, Reflect)] +#[reflect(Component)] +pub struct DecomposedPbrMaterial { + pub status: DecompositionStatus, + pub progress: f32, +} + +#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)] +pub struct PbrDecompositionLabel; + +// Copy of settings in render world +#[derive(Resource, Clone)] +pub struct RuntimeSettings(pub PbrDecompositionSettings); + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +pub struct NormalSettingsUniform { + spatial_sigma: f32, + color_sigma: f32, + confidence_threshold: f32, + _pad: f32, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +pub struct SyntheticViewSettingsUniform { + num_views: u32, + near_normal_angle_cos: f32, + sh_frame: u32, + _pad: u32, +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +pub struct MaterialSettingsUniform { + roughness_min: f32, + roughness_max: f32, + metallic_saturation_threshold: f32, + metallic_min_threshold: f32, +} + +pub fn extract_settings_to_render(mut render_commands: Commands, mut main_world: ResMut) { + let settings = main_world.resource::().clone(); + render_commands.insert_resource(RuntimeSettings(settings)); +} + +pub fn update_pbr_buffers( + render_device: Res, + gpu_planars: Res::GpuPlanarType>>, + mut normals_res: ResMut, + mut stats_res: ResMut, + mut mats_res: ResMut, +) { + for (asset_id, cloud) in gpu_planars.iter() { + let count = cloud.len(); + + normals_res + .asset_map + .entry(asset_id) + .or_insert_with(|| crate::pbr_decomposition::normal_estimation::GpuNormalBuffers::new(count, &render_device)); + + stats_res + .asset_map + .entry(asset_id) + .or_insert_with(|| GpuSyntheticViewsBuffers::new(count, &render_device)); + + mats_res + .asset_map + .entry(asset_id) + .or_insert_with(|| GpuMaterialBuffers::new(count, &render_device)); + } +} + +pub struct PbrDecompositionNode3d { + views: QueryState<( + &'static GaussianComputeViewBindGroup, + &'static ViewUniformOffset, + &'static PreviousViewUniformOffset, + )>, + gaussian_clouds: QueryState<( + &'static PlanarStorageBindGroup, + &'static PlanarGaussian3dHandle, + &'static DynamicUniformIndex, + Option<&'static GaussianMaterialOverrideBindGroup>, + )>, + initialized: bool, +} + +impl FromWorld for PbrDecompositionNode3d { + fn from_world(world: &mut World) -> Self { + Self { views: world.query(), gaussian_clouds: world.query(), initialized: false } + } +} + +impl Node for PbrDecompositionNode3d { + fn update(&mut self, world: &mut World) { + let pipeline = world.resource::(); + let pipeline_cache = world.resource::(); + if let CachedPipelineState::Ok(_) = pipeline_cache.get_compute_pipeline_state(pipeline.pipeline) { + self.initialized = true; + } + self.views.update_archetypes(world); + self.gaussian_clouds.update_archetypes(world); + } + + fn run(&self, _graph: &mut RenderGraphContext, render_context: &mut RenderContext, world: &World) -> Result<(), NodeRunError> { + if !self.initialized { return Ok(()); } + + let device = world.resource::(); + let pipeline_cache = world.resource::(); + let cloud_pipeline = world.resource::>(); + + let normals_pipeline = world.resource::(); + let synth_pipeline = world.resource::(); + let mats_pipeline = world.resource::(); + + let normals_res = world.resource::(); + let stats_res = world.resource::(); + let mats_res = world.resource::(); + + let rt_settings = world.resource::().0.clone(); + + let gaussian_uniforms = world.resource::(); + + let gpu_planars = world.resource::::GpuPlanarType>>(); + + for (view_bg, view_offset, prev_view_offset) in self.views.iter_manual(world) { + for (planar_bind_group, handle, cloud_uniform_index, material_override) in + self.gaussian_clouds.iter_manual(world) + { + let asset_id = handle.0.id(); + let Some(normals_gpu) = normals_res.asset_map.get(&asset_id) else { continue; }; + let Some(stats_gpu) = stats_res.asset_map.get(&asset_id) else { continue; }; + let Some(mats_gpu) = mats_res.asset_map.get(&asset_id) else { continue; }; + let Some(cloud_gpu) = gpu_planars.get(handle.0.id()) else { continue; }; + let gaussian_count = cloud_gpu.len() as u32; + let workgroups = gaussian_count.div_ceil(256).max(1); + + // Uniform buffers + let normal_settings = NormalSettingsUniform { + spatial_sigma: rt_settings.normal_spatial_sigma, + color_sigma: rt_settings.normal_color_sigma, + confidence_threshold: rt_settings.normal_confidence_threshold, + _pad: 0.0, + }; + let normals_uniform_buf = device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("normal_settings"), + contents: bytemuck::bytes_of(&normal_settings), + usage: BufferUsages::UNIFORM, + }); + + let synth_settings = SyntheticViewSettingsUniform { + num_views: rt_settings.num_synthetic_views, + near_normal_angle_cos: (rt_settings.view_near_normal_angle.to_radians()).cos(), + sh_frame: match rt_settings.sh_coordinate_frame { crate::pbr_decomposition::settings::SHCoordinateFrame::World => 0, _ => 1 }, + _pad: 0, + }; + let synth_uniform_buf = device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("synth_view_settings"), + contents: bytemuck::bytes_of(&synth_settings), + usage: BufferUsages::UNIFORM, + }); + + let mat_settings = MaterialSettingsUniform { + roughness_min: rt_settings.roughness_min, + roughness_max: rt_settings.roughness_max, + metallic_saturation_threshold: rt_settings.metallic_saturation_threshold, + metallic_min_threshold: rt_settings.metallic_min_threshold, + }; + let mats_uniform_buf = device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("material_settings"), + contents: bytemuck::bytes_of(&mat_settings), + usage: BufferUsages::UNIFORM, + }); + + // Bind groups for our pipelines (group 3 and 4) + let normals_bg = device.create_bind_group( + Some("pbr_normals_bg"), + &normals_pipeline.bind_group_layout, + &[BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { buffer: &normals_gpu.normals, offset: 0, size: None }), + }], + ); + let normals_settings_bg = device.create_bind_group( + Some("pbr_normals_settings_bg"), + &normals_pipeline.settings_layout, + &[BindGroupEntry { binding: 0, resource: normals_uniform_buf.as_entire_binding() }], + ); + + let synth_bg = device.create_bind_group( + Some("pbr_synth_bg"), + &synth_pipeline.bind_group_layout, + &[ + BindGroupEntry { binding: 0, resource: BindingResource::Buffer(BufferBinding { buffer: &normals_gpu.normals, offset: 0, size: None }) }, + BindGroupEntry { binding: 1, resource: BindingResource::Buffer(BufferBinding { buffer: &stats_gpu.stats, offset: 0, size: None }) }, + ], + ); + let synth_settings_bg = device.create_bind_group( + Some("pbr_synth_settings_bg"), + &synth_pipeline.settings_layout, + &[BindGroupEntry { binding: 0, resource: synth_uniform_buf.as_entire_binding() }], + ); + + let mats_bg = device.create_bind_group( + Some("pbr_mats_bg"), + &mats_pipeline.bind_group_layout, + &[ + BindGroupEntry { binding: 0, resource: BindingResource::Buffer(BufferBinding { buffer: &stats_gpu.stats, offset: 0, size: None }) }, + BindGroupEntry { binding: 1, resource: BindingResource::Buffer(BufferBinding { buffer: &mats_gpu.materials, offset: 0, size: None }) }, + ], + ); + let mats_settings_bg = device.create_bind_group( + Some("pbr_mats_settings_bg"), + &mats_pipeline.settings_layout, + &[BindGroupEntry { binding: 0, resource: mats_uniform_buf.as_entire_binding() }], + ); + + // Dispatch passes sequentially + debug!(gaussian_count, "dispatching PBR decomposition passes"); + let mut pass = render_context.command_encoder().begin_compute_pass(&ComputePassDescriptor::default()); + + // estimate_normals + if let Some(p) = pipeline_cache.get_compute_pipeline(normals_pipeline.pipeline) { + pass.set_pipeline(p); + pass.set_bind_group(0, &view_bg.value, &[view_offset.offset, prev_view_offset.offset]); + if let Some(gu) = &gaussian_uniforms.base_bind_group { pass.set_bind_group(1, gu, &[cloud_uniform_index.index()]); } + pass.set_bind_group(2, &planar_bind_group.bind_group, &[]); + pass.set_bind_group(3, &normals_bg, &[]); + pass.set_bind_group(4, &normals_settings_bg, &[]); + pass.dispatch_workgroups(workgroups, 1, 1); + } + + // evaluate_synthetic_views + if let Some(p) = pipeline_cache.get_compute_pipeline(synth_pipeline.pipeline) { + pass.set_pipeline(p); + pass.set_bind_group(0, &view_bg.value, &[view_offset.offset, prev_view_offset.offset]); + if let Some(gu) = &gaussian_uniforms.base_bind_group { pass.set_bind_group(1, gu, &[cloud_uniform_index.index()]); } + pass.set_bind_group(2, &planar_bind_group.bind_group, &[]); + pass.set_bind_group(3, &synth_bg, &[]); + pass.set_bind_group(4, &synth_settings_bg, &[]); + pass.dispatch_workgroups(workgroups, 1, 1); + } + + // estimate_material_properties + if let Some(p) = pipeline_cache.get_compute_pipeline(mats_pipeline.pipeline) { + pass.set_pipeline(p); + pass.set_bind_group(0, &view_bg.value, &[view_offset.offset, prev_view_offset.offset]); + if let Some(gu) = &gaussian_uniforms.base_bind_group { pass.set_bind_group(1, gu, &[cloud_uniform_index.index()]); } + pass.set_bind_group(2, &planar_bind_group.bind_group, &[]); + pass.set_bind_group(3, &mats_bg, &[]); + pass.set_bind_group(4, &mats_settings_bg, &[]); + let override_bind_group = material_override + .map(|bg| &bg.bind_group) + .unwrap_or(&cloud_pipeline.fallback_material_override_bind_group); + pass.set_bind_group(5, override_bind_group, &[]); + pass.dispatch_workgroups(workgroups, 1, 1); + } + drop(pass); + } + } + + Ok(()) + } +} diff --git a/src/pbr_decomposition/settings.rs b/src/pbr_decomposition/settings.rs new file mode 100644 index 00000000..3ff29117 --- /dev/null +++ b/src/pbr_decomposition/settings.rs @@ -0,0 +1,118 @@ +use bevy::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Reflect, Serialize, Deserialize)] +pub enum SHCoordinateFrame { + World, + Local, +} + +impl Default for SHCoordinateFrame { + fn default() -> Self { + Self::World + } +} + +#[derive(Resource, Reflect, Clone, Serialize, Deserialize)] +#[reflect(Resource)] +pub struct PbrDecompositionSettings { + pub use_gpu_hashed_grid: bool, + pub neighbor_search_radius: f32, + pub max_neighbors: u32, + pub hash_table_load_factor: f32, + + pub use_scale_axis_method: bool, + pub normal_spatial_sigma: f32, + pub normal_color_sigma: f32, + pub normal_confidence_threshold: f32, + + pub use_streaming_stats: bool, + pub num_synthetic_views: u32, + pub view_near_normal_angle: f32, + pub sh_coordinate_frame: SHCoordinateFrame, + pub topk_residuals: u32, + + pub enable_spectral_analysis: bool, + pub enable_chroma_filtering: bool, + pub specular_intensity_threshold: f32, + pub specular_hue_tolerance: f32, + + pub validated_highlight_min_count: u32, + pub metallic_saturation_threshold: f32, + pub metallic_min_threshold: f32, + pub roughness_min: f32, + pub roughness_max: f32, + pub ao_radius: f32, + pub ao_samples: u32, + + pub enable_material_propagation: bool, + pub propagation_radius: f32, + pub propagation_normal_threshold: f32, + pub propagation_blend_max: f32, + pub high_confidence_roughness_max: f32, + + pub enable_energy_validation: bool, + pub energy_error_threshold: f32, + pub ggx_lut_path: String, + + pub use_gbuffer_path: bool, + pub reconstruct_position_from_depth: bool, + pub enable_ssao: bool, + pub enable_intrinsic_ao: bool, + + pub cache_decomposed_materials: bool, + pub cache_path: String, +} + +impl Default for PbrDecompositionSettings { + fn default() -> Self { + Self { + use_gpu_hashed_grid: true, + neighbor_search_radius: 0.5, + max_neighbors: 64, + hash_table_load_factor: 0.5, + + use_scale_axis_method: true, + normal_spatial_sigma: 0.1, + normal_color_sigma: 0.2, + normal_confidence_threshold: 0.5, + + use_streaming_stats: true, + num_synthetic_views: 128, + view_near_normal_angle: 30.0, + sh_coordinate_frame: SHCoordinateFrame::World, + topk_residuals: 8, + + enable_spectral_analysis: true, + enable_chroma_filtering: true, + specular_intensity_threshold: 1.5, + specular_hue_tolerance: 0.1, + + validated_highlight_min_count: 3, + metallic_saturation_threshold: 0.15, + metallic_min_threshold: 0.1, + roughness_min: 0.089, + roughness_max: 1.0, + ao_radius: 0.5, + ao_samples: 16, + + enable_material_propagation: true, + propagation_radius: 1.0, + propagation_normal_threshold: 0.85, + propagation_blend_max: 0.5, + high_confidence_roughness_max: 0.3, + + enable_energy_validation: true, + energy_error_threshold: 0.1, + ggx_lut_path: "assets/textures/ggx_energy_lut.png".to_string(), + + use_gbuffer_path: true, + reconstruct_position_from_depth: true, + enable_ssao: true, + enable_intrinsic_ao: false, + + cache_decomposed_materials: true, + cache_path: "assets/cache/".to_string(), + } + } +} diff --git a/src/pbr_decomposition/shaders/material_separation.wgsl b/src/pbr_decomposition/shaders/material_separation.wgsl new file mode 100644 index 00000000..e702bdd2 --- /dev/null +++ b/src/pbr_decomposition/shaders/material_separation.wgsl @@ -0,0 +1,214 @@ +#define_import_path bevy_gaussian_splatting::pbr_decomposition::material_separation + +#import bevy_gaussian_splatting::bindings::{ + position_visibility, +} + +#ifdef PACKED + #ifdef PRECOMPUTE_COVARIANCE_3D + #import bevy_gaussian_splatting::packed::{ + get_position, + } + #else + #import bevy_gaussian_splatting::packed::{ + get_position, + } + #endif +#else ifdef BUFFER_STORAGE + #ifdef PRECOMPUTE_COVARIANCE_3D + #import bevy_gaussian_splatting::planar::{ + get_position, + } + #else + #import bevy_gaussian_splatting::planar::{ + get_position, + } + #endif +#else ifdef BUFFER_TEXTURE + #ifdef PRECOMPUTE_COVARIANCE_3D + #import bevy_gaussian_splatting::texture::{ + get_position, + } + #else + #import bevy_gaussian_splatting::texture::{ + get_position, + } + #endif +#endif + +struct StreamingStats { + mean_rgb: vec3, + count: u32, + + M2_rgb: vec3, + near_normal_count: u32, + + near_normal_mean: vec3, + topk_count: u32, + + topk_directions: array, 8>, + topk_intensities: array, + + residual_direction_sum: vec3, + residual_direction_M2: f32, + + _pad: array, +} + +struct PbrMaterialData { + base_color: vec3, + metallic: f32, + perceptual_roughness: f32, + reflectance: f32, + ambient_occlusion: f32, + _pad: f32, +} + +struct MaterialSettings { + roughness_min: f32, + roughness_max: f32, + metallic_saturation_threshold: f32, + metallic_min_threshold: f32, +} + +struct GaussianMaterialOverride { + base_color_factor: vec4, + bounds_min: vec4, + bounds_size: vec4, + flags: vec4, +} + +// Group 3: IO buffers for this pipeline +@group(3) @binding(0) var stats: array; +@group(3) @binding(1) var materials: array; + +// Group 4: settings +@group(4) @binding(0) var settings: MaterialSettings; + +// Group 5: material overrides +@group(5) @binding(0) var material_override: GaussianMaterialOverride; +@group(5) @binding(1) var base_color_texture: texture_2d; +@group(5) @binding(2) var base_color_sampler: sampler; + +const COLOR_EPSILON: f32 = 1e-6; + +fn compute_saturation_robust(color: vec3) -> f32 { + let c = clamp(color, vec3(0.0), vec3(1.0)); + + let max_c = max(c.r, max(c.g, c.b)); + + if (max_c < COLOR_EPSILON) { + return 0.0; + } + + let min_c = min(c.r, min(c.g, c.b)); + let delta = max_c - min_c; + + if (delta < COLOR_EPSILON) { + return 0.0; + } + + return delta / max_c; +} + +fn smoothstep_custom(edge0: f32, edge1: f32, x: f32) -> f32 { + let t = clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0); + return t * t * (3.0 - 2.0 * t); +} + +@compute @workgroup_size(256) +fn estimate_material_properties( + @builtin(global_invocation_id) global_id: vec3 +) { + let idx = global_id.x; + let gaussian_count = arrayLength(&stats); + if (idx >= gaussian_count) { return; } + + let stat = stats[idx]; + + var base_color = select( + stat.mean_rgb, + stat.near_normal_mean, + stat.near_normal_count > 8u + ); + + let base_color_factor = material_override.base_color_factor.rgb; + base_color *= base_color_factor; + + if (material_override.flags.x > 0u) { + let axis = material_override.flags.y; + let position = get_position(idx); + + var projected_position = vec2(0.0); + var bounds_projected_min = vec2(0.0); + var bounds_projected_size = vec2(1.0); + + if (axis == 0u) { + projected_position = position.xy; + bounds_projected_min = material_override.bounds_min.xy; + bounds_projected_size = material_override.bounds_size.xy; + } else if (axis == 1u) { + projected_position = vec2(position.x, position.z); + bounds_projected_min = vec2(material_override.bounds_min.x, material_override.bounds_min.z); + bounds_projected_size = vec2(material_override.bounds_size.x, material_override.bounds_size.z); + } else { + projected_position = position.yz; + bounds_projected_min = material_override.bounds_min.yz; + bounds_projected_size = material_override.bounds_size.yz; + } + + bounds_projected_size = max(bounds_projected_size, vec2(1e-3)); + var uv = (projected_position - bounds_projected_min) / bounds_projected_size; + uv = clamp(uv, vec2(0.0), vec2(1.0)); + + let texture_color = textureSampleLevel(base_color_texture, base_color_sampler, uv, 0.0).rgb; + base_color = texture_color * base_color; + } + + base_color = clamp(base_color, vec3(0.0), vec3(1.0)); + + var roughness = 0.5; + + if (stat.topk_count > 3u) { + let angular_std_dev = sqrt(stat.residual_direction_M2); + + roughness = mix( + settings.roughness_min, + settings.roughness_max, + saturate(angular_std_dev / 1.2) + ); + } else { + roughness = 0.8; + } + + let base_saturation = compute_saturation_robust(base_color); + + let base_energy = length(stat.near_normal_mean); + let spec_energy = length(stat.mean_rgb - stat.near_normal_mean); + let total_energy = base_energy + spec_energy; + let spec_ratio = spec_energy / max(0.0001, total_energy); + + var metallic = 0.0; + + if (base_saturation < settings.metallic_saturation_threshold && spec_ratio > 0.4) { + metallic = smoothstep_custom(0.4, 0.9, spec_ratio); + } else if (base_saturation > 0.3 && spec_ratio < 0.3) { + metallic = 0.0; + } else { + metallic = smoothstep_custom(0.5, 0.9, spec_ratio); + } + + metallic = select(0.0, metallic, metallic >= settings.metallic_min_threshold); + + let reflectance = 0.5; + let ao = 1.0; + + materials[idx] = PbrMaterialData( + base_color, + metallic, + roughness, + reflectance, + ao, + 0.0 + ); +} diff --git a/src/pbr_decomposition/shaders/normal_estimation.wgsl b/src/pbr_decomposition/shaders/normal_estimation.wgsl new file mode 100644 index 00000000..099e39a5 --- /dev/null +++ b/src/pbr_decomposition/shaders/normal_estimation.wgsl @@ -0,0 +1,94 @@ +#define_import_path bevy_gaussian_splatting::pbr_decomposition::normal_estimation + +#import bevy_gaussian_splatting::bindings::{ + position_visibility, + rotation, + scale_opacity, +} + +struct NormalData { + normal: vec3, + confidence: f32, +} + +struct NormalSettings { + spatial_sigma: f32, + color_sigma: f32, + confidence_threshold: f32, + _pad: f32, +} + +// Gaussian inputs come from the engine's planar storage bind group (group 2). +// Outputs and settings use dedicated groups (3 and 4) for this pipeline. +@group(3) @binding(0) var normals: array; + +@group(4) @binding(0) var settings: NormalSettings; + +fn quat_to_mat3(q: vec4) -> mat3x3 { + let qx = q.x; + let qy = q.y; + let qz = q.z; + let qw = q.w; + + let x2 = qx * qx; + let y2 = qy * qy; + let z2 = qz * qz; + let xy = qx * qy; + let xz = qx * qz; + let yz = qy * qz; + let wx = qw * qx; + let wy = qw * qy; + let wz = qw * qz; + + return mat3x3( + vec3(1.0 - 2.0 * (y2 + z2), 2.0 * (xy + wz), 2.0 * (xz - wy)), + vec3(2.0 * (xy - wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz + wx)), + vec3(2.0 * (xz + wy), 2.0 * (yz - wx), 1.0 - 2.0 * (x2 + y2)) + ); +} + +fn extract_normal_from_gaussian(rotation: vec4, scale: vec3) -> vec3 { + var min_scale_idx = 0u; + var min_scale = scale.x; + + if (scale.y < min_scale) { + min_scale = scale.y; + min_scale_idx = 1u; + } + + if (scale.z < min_scale) { + min_scale = scale.z; + min_scale_idx = 2u; + } + + let R = quat_to_mat3(rotation); + + var normal: vec3; + switch min_scale_idx { + case 0u: { normal = vec3(R[0][0], R[1][0], R[2][0]); } + case 1u: { normal = vec3(R[0][1], R[1][1], R[2][1]); } + case 2u: { normal = vec3(R[0][2], R[1][2], R[2][2]); } + default: { normal = vec3(0.0, 0.0, 1.0); } + } + + return normalize(normal); +} + +@compute @workgroup_size(256) +fn estimate_normals( + @builtin(global_invocation_id) global_id: vec3 +) { + let idx = global_id.x; + let gaussian_count = arrayLength(&position_visibility); + if (idx >= gaussian_count) { return; } + + let q = rotation[idx]; + let s = scale_opacity[idx].xyz; + + let normal = extract_normal_from_gaussian(q, s); + + // Minimal confidence until neighbor smoothing is wired with spatial hash + let confidence = 1.0; + + normals[idx] = NormalData(normal, confidence); +} diff --git a/src/pbr_decomposition/shaders/spatial_hash.wgsl b/src/pbr_decomposition/shaders/spatial_hash.wgsl new file mode 100644 index 00000000..7d0a0933 --- /dev/null +++ b/src/pbr_decomposition/shaders/spatial_hash.wgsl @@ -0,0 +1,107 @@ +#define_import_path bevy_gaussian_splatting::pbr_decomposition::spatial_hash + +struct SpatialHashConfig { + cell_size: f32, + table_size: u32, + gaussian_count: u32, + _pad: u32, +} + +struct GridCell { + start: u32, + count: u32, +} + +@group(0) @binding(0) var positions: array>; +@group(0) @binding(1) var cell_keys: array; +@group(0) @binding(2) var cell_indices: array; +@group(0) @binding(3) var cell_ranges: array; + +@group(1) @binding(0) var config: SpatialHashConfig; + +fn hash_position(pos: vec3) -> u32 { + let cell = vec3(floor(pos / config.cell_size)); + let p1 = u32(cell.x) * 73856093u; + let p2 = u32(cell.y) * 19349663u; + let p3 = u32(cell.z) * 83492791u; + return (p1 ^ p2 ^ p3) % config.table_size; +} + +@compute @workgroup_size(256) +fn compute_cell_keys( + @builtin(global_invocation_id) global_id: vec3 +) { + let idx = global_id.x; + if (idx >= config.gaussian_count) { return; } + + let pos = positions[idx]; + let hash = hash_position(pos); + + cell_keys[idx] = hash; + cell_indices[idx] = idx; +} + +@compute @workgroup_size(256) +fn build_cell_ranges( + @builtin(global_invocation_id) global_id: vec3 +) { + let idx = global_id.x; + if (idx >= config.gaussian_count) { return; } + + let key = cell_keys[idx]; + let prev_key = select(0xFFFFFFFFu, cell_keys[idx - 1u], idx > 0u); + + if (key != prev_key) { + cell_ranges[key].start = idx; + + if (prev_key != 0xFFFFFFFFu) { + let prev_start = cell_ranges[prev_key].start; + cell_ranges[prev_key].count = idx - prev_start; + } + } + + if (idx == config.gaussian_count - 1u) { + let start = cell_ranges[key].start; + cell_ranges[key].count = config.gaussian_count - start; + } +} + +fn query_neighbors_27( + query_pos: vec3, + max_neighbors: u32, + radius: f32, + neighbors: ptr> +) -> u32 { + let query_cell = vec3(floor(query_pos / config.cell_size)); + var count = 0u; + + for (var dx = -1; dx <= 1; dx++) { + for (var dy = -1; dy <= 1; dy++) { + for (var dz = -1; dz <= 1; dz++) { + let neighbor_cell = query_cell + vec3(dx, dy, dz); + + let cell_pos = vec3(neighbor_cell) * config.cell_size; + let hash = hash_position(cell_pos); + + let range = cell_ranges[hash]; + + for (var i = 0u; i < range.count; i++) { + if (count >= max_neighbors) { return count; } + + let candidate_idx = cell_indices[range.start + i]; + let candidate_pos = positions[candidate_idx]; + + let diff = query_pos - candidate_pos; + let dist_sq = dot(diff, diff); + + if (dist_sq <= radius * radius) { + (*neighbors)[count] = candidate_idx; + count++; + } + } + } + } + } + + return count; +} diff --git a/src/pbr_decomposition/shaders/synthetic_views.wgsl b/src/pbr_decomposition/shaders/synthetic_views.wgsl new file mode 100644 index 00000000..ebe8568b --- /dev/null +++ b/src/pbr_decomposition/shaders/synthetic_views.wgsl @@ -0,0 +1,194 @@ +#define_import_path bevy_gaussian_splatting::pbr_decomposition::synthetic_views + +#import bevy_gaussian_splatting::bindings::{ + spherical_harmonics, + rotation, +} +#import bevy_gaussian_splatting::spherical_harmonics::{ + spherical_harmonics_lookup, + srgb_to_linear, +} + +const PI: f32 = 3.14159265359; +const TOPK_RESIDUALS: u32 = 8u; + +struct StreamingStats { + mean_rgb: vec3, + count: u32, + + M2_rgb: vec3, + near_normal_count: u32, + + near_normal_mean: vec3, + topk_count: u32, + + topk_directions: array, 8>, + topk_intensities: array, + + residual_direction_sum: vec3, + residual_direction_M2: f32, + + _pad: array, +} + +struct SyntheticViewSettings { + num_views: u32, + near_normal_angle_cos: f32, + sh_frame: u32, + _pad: u32, +} + +struct NormalData { + normal: vec3, + confidence: f32, +} + +// Group 2: gaussian inputs via engine bindings (spherical_harmonics, rotation) +// Group 3: this pipeline IO: normals (read) + stats (read_write) +// Group 4: settings +@group(3) @binding(0) var normals: array; +@group(3) @binding(1) var stats: array; + +@group(4) @binding(0) var settings: SyntheticViewSettings; + +const SH_FRAME_WORLD: u32 = 0u; +const SH_FRAME_LOCAL: u32 = 1u; + +fn fibonacci_sphere(i: u32, n: u32) -> vec3 { + let phi = PI * (sqrt(5.0) - 1.0); + let y = 1.0 - (f32(i) / f32(n - 1u)) * 2.0; + let radius = sqrt(1.0 - y * y); + let theta = phi * f32(i); + + return vec3( + cos(theta) * radius, + y, + sin(theta) * radius + ); +} + +fn quat_to_mat3_transpose(q: vec4) -> mat3x3 { + let qx = q.x; + let qy = q.y; + let qz = q.z; + let qw = q.w; + + let x2 = qx * qx; + let y2 = qy * qy; + let z2 = qz * qz; + let xy = qx * qy; + let xz = qx * qz; + let yz = qy * qz; + let wx = qw * qx; + let wy = qw * qy; + let wz = qw * qz; + + return transpose(mat3x3( + vec3(1.0 - 2.0 * (y2 + z2), 2.0 * (xy + wz), 2.0 * (xz - wy)), + vec3(2.0 * (xy - wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz + wx)), + vec3(2.0 * (xz + wy), 2.0 * (yz - wx), 1.0 - 2.0 * (x2 + y2)) + )); +} + +fn update_topk_residuals( + stats_ptr: ptr, + direction: vec3, + intensity: f32 +) { + if ((*stats_ptr).topk_count < TOPK_RESIDUALS) { + let idx = (*stats_ptr).topk_count; + (*stats_ptr).topk_directions[idx] = direction; + (*stats_ptr).topk_intensities[idx] = intensity; + (*stats_ptr).topk_count++; + } else { + var min_idx = 0u; + var min_val = (*stats_ptr).topk_intensities[0]; + for (var i = 1u; i < TOPK_RESIDUALS; i++) { + if ((*stats_ptr).topk_intensities[i] < min_val) { + min_val = (*stats_ptr).topk_intensities[i]; + min_idx = i; + } + } + if (intensity > min_val) { + (*stats_ptr).topk_directions[min_idx] = direction; + (*stats_ptr).topk_intensities[min_idx] = intensity; + } + } +} + +@compute @workgroup_size(256) +fn evaluate_synthetic_views( + @builtin(global_invocation_id) global_id: vec3 +) { + let idx = global_id.x; + let gaussian_count = arrayLength(&normals); + if (idx >= gaussian_count) { return; } + + let sh = spherical_harmonics[idx]; + let normal_data = normals[idx]; + let normal = normal_data.normal; + let rotation = rotation[idx]; + + var local_stats: StreamingStats; + local_stats.count = 0u; + local_stats.mean_rgb = vec3(0.0); + local_stats.M2_rgb = vec3(0.0); + local_stats.near_normal_count = 0u; + local_stats.near_normal_mean = vec3(0.0); + local_stats.topk_count = 0u; + local_stats.residual_direction_sum = vec3(0.0); + local_stats.residual_direction_M2 = 0.0; + + for (var i = 0u; i < settings.num_views; i++) { + let view_dir_world = fibonacci_sphere(i, settings.num_views); + + var view_dir_eval = view_dir_world; + if (settings.sh_frame == SH_FRAME_LOCAL) { + let R_inv = quat_to_mat3_transpose(rotation); + view_dir_eval = R_inv * view_dir_world; + } + + let color_srgb = spherical_harmonics_lookup(view_dir_eval, sh); + let color = max(vec3(0.0), srgb_to_linear(color_srgb)); + + local_stats.count += 1u; + let delta = color - local_stats.mean_rgb; + local_stats.mean_rgb += delta / f32(local_stats.count); + let delta2 = color - local_stats.mean_rgb; + local_stats.M2_rgb += delta * delta2; + + let ndotv = dot(normal, view_dir_world); + if (ndotv > settings.near_normal_angle_cos) { + local_stats.near_normal_count += 1u; + let near_delta = color - local_stats.near_normal_mean; + local_stats.near_normal_mean += near_delta / f32(local_stats.near_normal_count); + } + + if (local_stats.near_normal_count > 0u) { + let base_estimate = local_stats.near_normal_mean; + let lambertian_prediction = base_estimate * max(0.0, ndotv); + + let residual = color - lambertian_prediction; + let residual_intensity = length(residual); + + if (residual_intensity > 0.01) { + update_topk_residuals(&local_stats, view_dir_world, residual_intensity); + local_stats.residual_direction_sum += view_dir_world; + } + } + } + + if (local_stats.topk_count > 3u) { + let mean_direction = normalize(local_stats.residual_direction_sum); + + var angular_variance = 0.0; + for (var i = 0u; i < local_stats.topk_count; i++) { + let dir = local_stats.topk_directions[i]; + let angle = acos(clamp(dot(dir, mean_direction), -1.0, 1.0)); + angular_variance += angle * angle; + } + local_stats.residual_direction_M2 = angular_variance / f32(local_stats.topk_count); + } + + stats[idx] = local_stats; +} diff --git a/src/pbr_decomposition/spatial_hash.rs b/src/pbr_decomposition/spatial_hash.rs new file mode 100644 index 00000000..70a559f5 --- /dev/null +++ b/src/pbr_decomposition/spatial_hash.rs @@ -0,0 +1,5 @@ +use bevy::prelude::*; + +pub mod pipeline; + +pub use pipeline::*; diff --git a/src/pbr_decomposition/spatial_hash/pipeline.rs b/src/pbr_decomposition/spatial_hash/pipeline.rs new file mode 100644 index 00000000..024f34a3 --- /dev/null +++ b/src/pbr_decomposition/spatial_hash/pipeline.rs @@ -0,0 +1,189 @@ +use bevy::{ + asset::{load_internal_asset, uuid_handle}, + prelude::*, + render::{ + render_resource::*, + renderer::RenderDevice, + }, +}; +use std::collections::HashMap; +use bevy::asset::AssetId; + +use crate::pbr_decomposition::types::*; + +const SPATIAL_HASH_SHADER_HANDLE: Handle = + uuid_handle!("a1b2c3d4-e5f6-7890-abcd-ef1234567890"); + +#[derive(Resource)] +pub struct SpatialHashBuffers { + pub asset_map: HashMap, GpuSpatialHashBuffers>, +} + +impl Default for SpatialHashBuffers { + fn default() -> Self { + SpatialHashBuffers { + asset_map: HashMap::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct GpuSpatialHashBuffers { + pub cell_keys: Buffer, + pub cell_indices: Buffer, + pub cell_ranges: Buffer, + pub config: SpatialHashConfig, +} + +impl GpuSpatialHashBuffers { + pub fn new(config: SpatialHashConfig, render_device: &RenderDevice) -> Self { + let count = config.gaussian_count as usize; + let table_size = config.table_size as usize; + + let cell_keys = render_device.create_buffer(&BufferDescriptor { + label: Some("spatial hash cell keys"), + size: (count * std::mem::size_of::()) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + let cell_indices = render_device.create_buffer(&BufferDescriptor { + label: Some("spatial hash cell indices"), + size: (count * std::mem::size_of::()) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + let cell_ranges = render_device.create_buffer(&BufferDescriptor { + label: Some("spatial hash cell ranges"), + size: (table_size * std::mem::size_of::()) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + Self { + cell_keys, + cell_indices, + cell_ranges, + config, + } + } +} + +#[derive(Resource)] +pub struct SpatialHashPipeline { + pub compute_keys_pipeline: CachedComputePipelineId, + pub build_ranges_pipeline: CachedComputePipelineId, + pub bind_group_layout: BindGroupLayout, + pub config_layout: BindGroupLayout, +} + +impl FromWorld for SpatialHashPipeline { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + let bind_group_layout = render_device.create_bind_group_layout( + Some("spatial_hash_bind_group_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + ); + + let config_layout = render_device.create_bind_group_layout( + Some("spatial_hash_config_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + ); + + let pipeline_cache = world.resource::(); + + let compute_keys_pipeline = pipeline_cache.queue_compute_pipeline( + ComputePipelineDescriptor { + label: Some("spatial_hash_compute_keys".into()), + layout: vec![bind_group_layout.clone(), config_layout.clone()], + push_constant_ranges: vec![], + shader: SPATIAL_HASH_SHADER_HANDLE, + shader_defs: vec![], + entry_point: Some("compute_cell_keys".into()), + zero_initialize_workgroup_memory: false, + }, + ); + + let build_ranges_pipeline = pipeline_cache.queue_compute_pipeline( + ComputePipelineDescriptor { + label: Some("spatial_hash_build_ranges".into()), + layout: vec![bind_group_layout.clone(), config_layout.clone()], + push_constant_ranges: vec![], + shader: SPATIAL_HASH_SHADER_HANDLE, + shader_defs: vec![], + entry_point: Some("build_cell_ranges".into()), + zero_initialize_workgroup_memory: false, + }, + ); + + Self { + compute_keys_pipeline, + build_ranges_pipeline, + bind_group_layout, + config_layout, + } + } +} + +pub fn load_spatial_hash_shader(app: &mut App) { + load_internal_asset!( + app, + SPATIAL_HASH_SHADER_HANDLE, + "../shaders/spatial_hash.wgsl", + Shader::from_wgsl + ); +} diff --git a/src/pbr_decomposition/synthetic_views.rs b/src/pbr_decomposition/synthetic_views.rs new file mode 100644 index 00000000..70a559f5 --- /dev/null +++ b/src/pbr_decomposition/synthetic_views.rs @@ -0,0 +1,5 @@ +use bevy::prelude::*; + +pub mod pipeline; + +pub use pipeline::*; diff --git a/src/pbr_decomposition/synthetic_views/pipeline.rs b/src/pbr_decomposition/synthetic_views/pipeline.rs new file mode 100644 index 00000000..92ff13d1 --- /dev/null +++ b/src/pbr_decomposition/synthetic_views/pipeline.rs @@ -0,0 +1,138 @@ +use bevy::{ + asset::{load_internal_asset, uuid_handle}, + prelude::*, + render::{ + render_resource::*, + renderer::RenderDevice, + }, +}; +use std::collections::HashMap; +use bevy::asset::AssetId; +use crate::gaussian::formats::planar_3d::PlanarGaussian3d; + +use crate::pbr_decomposition::types::*; +use crate::render::{CloudPipeline, CloudPipelineKey, shader_defs}; +use crate::gaussian::formats::planar_3d::Gaussian3d; + +const SYNTHETIC_VIEWS_SHADER_HANDLE: Handle = + uuid_handle!("c3d4e5f6-a7b8-9012-cdef-123456789012"); + +#[derive(Resource)] +pub struct SyntheticViewsBuffers { + pub asset_map: HashMap, GpuSyntheticViewsBuffers>, +} + +impl Default for SyntheticViewsBuffers { + fn default() -> Self { + SyntheticViewsBuffers { + asset_map: HashMap::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct GpuSyntheticViewsBuffers { + pub stats: Buffer, +} + +impl GpuSyntheticViewsBuffers { + pub fn new(count: usize, render_device: &RenderDevice) -> Self { + let stats = render_device.create_buffer(&BufferDescriptor { + label: Some("streaming_stats"), + size: (count as u64) * StreamingStats::min_size().get(), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + Self { stats } + } +} + +#[derive(Resource)] +pub struct SyntheticViewsPipeline { + pub pipeline: CachedComputePipelineId, + pub bind_group_layout: BindGroupLayout, + pub settings_layout: BindGroupLayout, +} + +impl FromWorld for SyntheticViewsPipeline { + fn from_world(world: &mut World) -> Self { + let render_device = world.resource::(); + + // Group 3: normals (read) + stats (read_write) + let bind_group_layout = render_device.create_bind_group_layout( + Some("synthetic_views_bind_group_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: Some(NormalData::min_size()), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: Some(StreamingStats::min_size()), + }, + count: None, + }, + ], + ); + + let settings_layout = render_device.create_bind_group_layout( + Some("synthetic_views_settings_layout"), + &[BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }], + ); + + let pipeline_cache = world.resource::(); + let cloud_pipeline = world.resource::>(); + let shader_defs = shader_defs(CloudPipelineKey::default()); + + let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("synthetic_views".into()), + layout: vec![ + cloud_pipeline.compute_view_layout.clone(), + cloud_pipeline.gaussian_uniform_layout.clone(), + cloud_pipeline.gaussian_cloud_layout.clone(), + bind_group_layout.clone(), + settings_layout.clone(), + ], + push_constant_ranges: vec![], + shader: SYNTHETIC_VIEWS_SHADER_HANDLE, + shader_defs, + entry_point: Some("evaluate_synthetic_views".into()), + zero_initialize_workgroup_memory: false, + }); + + Self { + pipeline, + bind_group_layout, + settings_layout, + } + } +} + +pub fn load_synthetic_views_shader(app: &mut App) { + load_internal_asset!( + app, + SYNTHETIC_VIEWS_SHADER_HANDLE, + "../shaders/synthetic_views.wgsl", + Shader::from_wgsl + ); +} diff --git a/src/pbr_decomposition/types.rs b/src/pbr_decomposition/types.rs new file mode 100644 index 00000000..a98e5bce --- /dev/null +++ b/src/pbr_decomposition/types.rs @@ -0,0 +1,192 @@ +use bevy::prelude::*; +use bevy::render::render_resource::ShaderType; +use bytemuck::{Pod, Zeroable}; +use serde::{Deserialize, Serialize}; + +#[derive( + Clone, + Debug, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct PbrMaterialData { + pub base_color: [f32; 3], + pub metallic: f32, + pub perceptual_roughness: f32, + pub reflectance: f32, + pub ambient_occlusion: f32, + pub _pad: f32, +} + +impl Default for PbrMaterialData { + fn default() -> Self { + Self { + base_color: [0.5, 0.5, 0.5], + metallic: 0.0, + perceptual_roughness: 0.5, + reflectance: 0.5, + ambient_occlusion: 1.0, + _pad: 0.0, + } + } +} + +#[derive( + Clone, + Debug, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct NormalData { + pub normal: [f32; 3], + pub confidence: f32, +} + +impl Default for NormalData { + fn default() -> Self { + Self { + normal: [0.0, 0.0, 1.0], + confidence: 0.0, + } + } +} + +#[derive(Clone, Debug, Copy, PartialEq, Reflect, ShaderType)] +#[repr(C)] +pub struct StreamingStats { + pub mean_rgb: Vec3, + pub count: u32, + + pub M2_rgb: Vec3, + pub near_normal_count: u32, + + pub near_normal_mean: Vec3, + pub topk_count: u32, + + pub topk_directions: [Vec3; 8], + pub topk_intensities: [f32; 8], + + pub residual_direction_sum: Vec3, + pub residual_direction_M2: f32, + + pub _pad: [f32; 3], +} + +impl Default for StreamingStats { + fn default() -> Self { + Self { + mean_rgb: Vec3::ZERO, + count: 0, + M2_rgb: Vec3::ZERO, + near_normal_count: 0, + near_normal_mean: Vec3::ZERO, + topk_count: 0, + topk_directions: [Vec3::ZERO; 8], + topk_intensities: [0.0; 8], + residual_direction_sum: Vec3::ZERO, + residual_direction_M2: 0.0, + _pad: [0.0; 3], + } + } +} + +#[derive( + Clone, + Debug, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct SpatialHashConfig { + pub cell_size: f32, + pub table_size: u32, + pub gaussian_count: u32, + pub _pad: u32, +} + +impl Default for SpatialHashConfig { + fn default() -> Self { + Self { + cell_size: 0.5, + table_size: 1024, + gaussian_count: 0, + _pad: 0, + } + } +} + +impl SpatialHashConfig { + pub fn from_point_cloud(positions: &[[f32; 3]], target_neighbors: u32) -> Self { + if positions.is_empty() { + return Self::default(); + } + + let mut min = [f32::MAX; 3]; + let mut max = [f32::MIN; 3]; + + for pos in positions { + for i in 0..3 { + min[i] = min[i].min(pos[i]); + max[i] = max[i].max(pos[i]); + } + } + + let volume = (max[0] - min[0]) * (max[1] - min[1]) * (max[2] - min[2]); + let density = positions.len() as f32 / volume.max(1e-6); + + let cell_size = (target_neighbors as f32 / density).cbrt().max(0.01); + + let table_size = (positions.len() * 2).next_power_of_two(); + + Self { + cell_size, + table_size: table_size as u32, + gaussian_count: positions.len() as u32, + _pad: 0, + } + } +} + +#[derive( + Clone, + Debug, + Copy, + PartialEq, + Reflect, + ShaderType, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct GridCell { + pub start: u32, + pub count: u32, +} + +impl Default for GridCell { + fn default() -> Self { + Self { start: 0, count: 0 } + } +} diff --git a/src/render/bindings.wgsl b/src/render/bindings.wgsl index 481caf25..0bae8113 100644 --- a/src/render/bindings.wgsl +++ b/src/render/bindings.wgsl @@ -2,9 +2,8 @@ #import bevy_pbr::prepass_bindings::PreviousViewUniforms #import bevy_render::globals::Globals -#import bevy_render::view::View +#import bevy_pbr::mesh_view_bindings::view -@group(0) @binding(0) var view: View; @group(0) @binding(1) var globals: Globals; @group(0) @binding(2) var previous_view_uniforms: PreviousViewUniforms; diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index dc7f8b14..5fd3e0e5 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -1,8 +1,4 @@ -#import bevy_gaussian_splatting::bindings::{ - view, - gaussian_uniforms, - Entry, -} +#import bevy_gaussian_splatting::bindings::{gaussian_uniforms, Entry} #import bevy_gaussian_splatting::classification::class_to_rgb #import bevy_gaussian_splatting::depth::depth_to_rgb #import bevy_gaussian_splatting::optical_flow::{ @@ -18,6 +14,12 @@ in_frustum, } +#import bevy_pbr::{ + pbr_functions::{apply_pbr_lighting, calculate_view, main_pass_post_lighting_processing}, + pbr_types, +}; +#import bevy_pbr::mesh_view_bindings as view_bindings; + #ifdef GAUSSIAN_2D #import bevy_gaussian_splatting::gaussian_2d::{ compute_cov2d_surfel, @@ -123,6 +125,23 @@ } #endif +struct NormalData { + normal: vec3, + confidence: f32, +} + +struct PbrMaterialData { + base_color: vec3, + metallic: f32, + perceptual_roughness: f32, + reflectance: f32, + ambient_occlusion: f32, + _pad: f32, +} + +@group(4) @binding(0) var decomposed_normals: array; +@group(4) @binding(1) var pbr_materials: array; + #ifdef WEBGL2 struct GaussianVertexOutput { @builtin(position) position: vec4, @@ -141,6 +160,7 @@ @location(2) conic: vec3, @location(3) major_minor: vec2, #endif + @location(7) gaussian_index: u32, }; #else struct GaussianVertexOutput { @@ -160,6 +180,7 @@ @location(2) @interpolate(flat) conic: vec3, @location(3) @interpolate(linear) major_minor: vec2, #endif + @location(7) @interpolate(flat) gaussian_index: u32, }; #endif @@ -313,7 +334,7 @@ fn vs_points( // TODO: RASTERIZE_ACCELERATION #ifdef RASTERIZE_CLASSIFICATION - let ray_direction_world = normalize(transformed_position - view.world_position); + let ray_direction_world = normalize(transformed_position - view_bindings::view.world_position); let ray_direction_local = world_to_local_direction(ray_direction_world, gaussian_uniforms.transform); #ifdef GAUSSIAN_3D_STRUCTURE @@ -334,7 +355,7 @@ fn vs_points( let min_position = (gaussian_uniforms.transform * first_position).xyz; let max_position = (gaussian_uniforms.transform * last_position).xyz; - let camera_position = view.world_position; + let camera_position = view_bindings::view.world_position; let min_distance = length(min_position - camera_position); let max_distance = length(max_position - camera_position); @@ -357,7 +378,7 @@ fn vs_points( let L = T * S * R; let local_normal = vec4(L[2], 0.0); - let world_normal = view.view_from_world * local_normal; + let world_normal = view_bindings::view.view_from_world * local_normal; let t = normalize(world_normal); @@ -405,7 +426,7 @@ fn vs_points( rgb = base_color * scaled_mag; #else ifdef RASTERIZE_COLOR // TODO: verify color benefit for ray_direction computed at quad verticies instead of gaussian center (same as current complexity) - let ray_direction_world = normalize(transformed_position - view.world_position); + let ray_direction_world = normalize(transformed_position - view_bindings::view.world_position); let ray_direction_local = world_to_local_direction(ray_direction_world, gaussian_uniforms.transform); #ifdef GAUSSIAN_3D_STRUCTURE @@ -431,6 +452,7 @@ fn vs_points( projected_position.xy + bb.xy, projected_position.zw, ); + output.gaussian_index = splat_index; return output; } @@ -443,7 +465,7 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { let mean_2d = input.mean_2d; let aspect = vec2( 1.0, - view.viewport.z / view.viewport.w, + view_bindings::view.viewport.z / view_bindings::view.viewport.w, ); let pixel_coord = input.uv * radius * aspect + mean_2d; @@ -496,10 +518,65 @@ fn fs_main(input: GaussianVertexOutput) -> @location(0) vec4 { let alpha = min(exp(power) * input.color.a, 0.999); - // TODO: round alpha to terminate depth test? +#ifdef HIGHLIGHT_SELECTED + if (get_visibility(input.gaussian_index) > 0.5) { + return vec4(input.color.rgb * alpha, alpha); + } +#endif + +#ifdef RASTERIZE_COLOR + if (alpha <= 0.0) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + + if (input.gaussian_index >= arrayLength(&pbr_materials)) { + return vec4(input.color.rgb * alpha, alpha); + } + + let material = pbr_materials[input.gaussian_index]; + let normal_data = decomposed_normals[input.gaussian_index]; + let clamped_metallic = clamp(material.metallic, 0.0, 1.0); + let clamped_roughness = clamp(material.perceptual_roughness, 0.089, 1.0); + let clamped_reflectance = clamp(material.reflectance, 0.0, 1.0); + let clamped_occlusion = clamp(material.ambient_occlusion, 0.0, 1.0); + + var world_normal = normalize(normal_data.normal); + let has_nan = any(world_normal != world_normal); + if (normal_data.confidence <= 0.0 || has_nan) { + world_normal = vec3(0.0, 0.0, 1.0); + } + + let local_position = get_position(input.gaussian_index); + let world_position = (gaussian_uniforms.transform * vec4(local_position, 1.0)).xyz; + let is_orthographic = view_bindings::view.clip_from_view[3].w == 1.0; + + var pbr_input = pbr_types::pbr_input_new(); + pbr_input.material.base_color = vec4(material.base_color, 1.0); + pbr_input.material.metallic = clamped_metallic; + pbr_input.material.perceptual_roughness = clamped_roughness; + pbr_input.material.reflectance = vec3(clamped_reflectance); + pbr_input.material.flags = + pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND | + pbr_types::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT | + pbr_types::STANDARD_MATERIAL_FLAGS_FOG_ENABLED_BIT; + pbr_input.material.deferred_lighting_pass_id = 1u; + pbr_input.diffuse_occlusion = vec3(clamped_occlusion); + pbr_input.specular_occlusion = clamped_occlusion; + pbr_input.world_position = vec4(world_position, 1.0); + pbr_input.is_orthographic = is_orthographic; + pbr_input.V = calculate_view(pbr_input.world_position, is_orthographic); + pbr_input.world_normal = world_normal; + pbr_input.N = world_normal; + + var lit_color = apply_pbr_lighting(pbr_input); + lit_color = main_pass_post_lighting_processing(pbr_input, lit_color); + + return vec4(lit_color.rgb * alpha, alpha); +#else return vec4( input.color.rgb * alpha, alpha, ); +#endif } diff --git a/src/render/helpers.wgsl b/src/render/helpers.wgsl index 0fb1fc20..8e141da4 100644 --- a/src/render/helpers.wgsl +++ b/src/render/helpers.wgsl @@ -1,9 +1,7 @@ #define_import_path bevy_gaussian_splatting::helpers -#import bevy_gaussian_splatting::bindings::{ - view, - gaussian_uniforms, -} +#import bevy_gaussian_splatting::bindings::gaussian_uniforms +#import bevy_pbr::mesh_view_bindings as view_bindings fn cov2d( position: vec3, @@ -15,11 +13,11 @@ fn cov2d( cov3d[2], cov3d[4], cov3d[5], ); - var t = view.view_from_world * vec4(position, 1.0); + var t = view_bindings::view.view_from_world * vec4(position, 1.0); let focal = vec2( - view.clip_from_view[0].x * view.viewport.z, - view.clip_from_view[1].y * view.viewport.w, + view_bindings::view.clip_from_view[0].x * view_bindings::view.viewport.z, + view_bindings::view.clip_from_view[1].y * view_bindings::view.viewport.w, ); let s = 1.0 / (t.z * t.z); @@ -31,9 +29,9 @@ fn cov2d( let W = transpose( mat3x3( - view.view_from_world[0].xyz, - view.view_from_world[1].xyz, - view.view_from_world[2].xyz, + view_bindings::view.view_from_world[0].xyz, + view_bindings::view.view_from_world[1].xyz, + view_bindings::view.view_from_world[2].xyz, ) ); @@ -69,7 +67,7 @@ fn get_bounding_box_clip( #ifdef USE_AABB let radius_px = cutoff * max(x_axis_length, y_axis_length); let radius_ndc = vec2( - radius_px / view.viewport.zw, + radius_px / view_bindings::view.viewport.zw, ); return vec4( @@ -109,7 +107,7 @@ fn get_bounding_box_clip( let scaled_vertex = direction * bounds; let rotated_vertex = scaled_vertex * rotation_matrix; - let scaling_factor = 1.0 / view.viewport.zw; + let scaling_factor = 1.0 / view_bindings::view.viewport.zw; let ndc_vertex = rotated_vertex * scaling_factor; return vec4( @@ -121,13 +119,13 @@ fn get_bounding_box_clip( fn intrinsic_matrix() -> mat3x4 { let focal = vec2( - view.clip_from_view[0].x * view.viewport.z / 2.0, - view.clip_from_view[1].y * view.viewport.w / 2.0, + view_bindings::view.clip_from_view[0].x * view_bindings::view.viewport.z / 2.0, + view_bindings::view.clip_from_view[1].y * view_bindings::view.viewport.w / 2.0, ); let Ks = mat3x4( - vec4(focal.x, 0.0, 0.0, (view.viewport.z - 1.0) / 2.0), - vec4(0.0, focal.y, 0.0, (view.viewport.w - 1.0) / 2.0), + vec4(focal.x, 0.0, 0.0, (view_bindings::view.viewport.z - 1.0) / 2.0), + vec4(0.0, focal.y, 0.0, (view_bindings::view.viewport.w - 1.0) / 2.0), vec4(0.0, 0.0, 0.0, 1.0) ); diff --git a/src/render/mod.rs b/src/render/mod.rs index cfc49610..97e4bb1e 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -1,32 +1,46 @@ #![allow(dead_code)] // ShaderType derives emit unused check helpers -use std::{hash::Hash, num::NonZero}; +use std::{any::TypeId, hash::Hash, num::NonZero}; +use bevy::render::render_resource::TextureFormat; +use bevy::shader::ShaderDefVal; use bevy::{ - asset::{load_internal_asset, uuid_handle, AssetEvent, AssetId}, camera::primitives::Aabb, core_pipeline::{ + asset::{AssetEvent, AssetId, load_internal_asset, uuid_handle}, + camera::primitives::Aabb, + core_pipeline::{ core_3d::Transparent3d, - prepass::{ - MotionVectorPrepass, PreviousViewData, PreviousViewUniformOffset, PreviousViewUniforms, - }, - }, ecs::{ + prepass::{PreviousViewData, PreviousViewUniforms}, + }, + ecs::{ query::ROQueryItem, - system::{lifetimeless::*, SystemParamItem}, - }, pbr::PrepassViewBindGroup, prelude::*, render::{ - extract_component::{ComponentUniforms, DynamicUniformIndex, UniformComponentPlugin}, globals::{GlobalsBuffer, GlobalsUniform}, render_asset::RenderAssets, render_phase::{ - AddRenderCommand, DrawFunctions, PhaseItem, PhaseItemExtraIndex, RenderCommand, - RenderCommandResult, SetItemPipeline, TrackedRenderPass, ViewSortedRenderPhases, - }, render_resource::*, renderer::RenderDevice, sync_world::RenderEntity, view::{ - ExtractedView, RenderVisibilityRanges, RenderVisibleEntities, ViewUniform, ViewUniformOffset, ViewUniforms, VISIBILITY_RANGES_STORAGE_BUFFER_COUNT - }, Extract, Render, RenderApp, RenderSystems - } + system::{SystemParamItem, lifetimeless::*}, + }, + prelude::*, + render::{ + Extract, Render, RenderApp, RenderSystems, + extract_component::{ComponentUniforms, DynamicUniformIndex, UniformComponentPlugin}, + globals::{GlobalsBuffer, GlobalsUniform}, + render_asset::RenderAssets, render_phase::{ + AddRenderCommand, DrawFunctions, PhaseItem, PhaseItemExtraIndex, RenderCommand, + RenderCommandResult, SetItemPipeline, TrackedRenderPass, ViewSortedRenderPhases, + }, + render_resource::*, renderer::RenderDevice, sync_world::RenderEntity, + view::{ + ExtractedView, RenderVisibilityRanges, RenderVisibleEntities, + VISIBILITY_RANGES_STORAGE_BUFFER_COUNT, ViewUniform, ViewUniformOffset, ViewUniforms, + }, + }, }; -use bevy::shader::ShaderDefVal; use bevy_interleave::prelude::*; -use bevy::render::render_resource::TextureFormat; +use bytemuck::{bytes_of, Pod, Zeroable}; +use bevy::core_pipeline::{oit::OrderIndependentTransparencySettings, prepass::ViewPrepassTextures}; +use bevy_pbr::{MeshPipelineViewLayoutKey, MeshPipelineViewLayouts, SetMeshViewBindGroup}; +use bevy::render::texture::FallbackImage; use crate::{ camera::GaussianCamera, gaussian::{ cloud::CloudVisibilityClass, + formats::planar_3d::{Gaussian3d, PlanarGaussian3dHandle}, interface::CommonCloud, settings::{CloudSettings, DrawMode, GaussianMode, RasterizeMode}, }, @@ -35,6 +49,12 @@ use crate::{ spherindrical_harmonics::{SH_4D_COEFF_COUNT, SH_4D_DEGREE_TIME}, }, morph::MorphPlugin, + pbr_decomposition::{ + material_separation::MaterialSeparationBuffers, + normal_estimation::NormalEstimationBuffers, + types::{NormalData, PbrMaterialData}, + }, + material::gaussian_material::{GaussianMaterial, GaussianMaterialBounds, GaussianMaterialHandle, RenderGaussianMaterials}, sort::{GpuSortedEntry, SortEntry, SortPlugin, SortTrigger, SortedEntriesHandle}, }; @@ -62,6 +82,15 @@ const TEXTURE_SHADER_HANDLE: Handle = uuid_handle!("500e2ebf-51a8-402e-9 const TRANSFORM_SHADER_HANDLE: Handle = uuid_handle!("648516b2-87cc-4937-ae1c-d986952e9fa7"); +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct GaussianMaterialOverrideUniform { + base_color_factor: [f32; 4], + bounds_min: [f32; 4], + bounds_size: [f32; 4], + flags: [u32; 4], +} + // TODO: consider refactor to bind via bevy's mesh (dynamic vertex planes) + shared batching/instancing/preprocessing // utilize RawBufferVec for gaussian data? pub struct RenderPipelinePlugin { @@ -107,12 +136,20 @@ where refresh_planar_storage_bind_groups:: .in_set(RenderSystems::PrepareBindGroups), queue_gaussian_bind_group::.in_set(RenderSystems::PrepareBindGroups), - queue_gaussian_view_bind_groups::.in_set(RenderSystems::PrepareBindGroups), + queue_gaussian_view_bind_groups:: + .in_set(RenderSystems::PrepareBindGroups), queue_gaussian_compute_view_bind_groups:: .in_set(RenderSystems::PrepareBindGroups), queue_gaussians::.in_set(RenderSystems::Queue), ), ); + + if TypeId::of::() == TypeId::of::() { + render_app.add_systems( + Render, + queue_gaussian_pbr_bind_group_3d.in_set(RenderSystems::PrepareBindGroups), + ); + } } // TODO: refactor common resources into a common plugin @@ -354,12 +391,14 @@ fn queue_gaussians( &GaussianCamera, &RenderVisibleEntities, Option<&Msaa>, + Option<&ViewPrepassTextures>, + Has, )>, gaussian_splatting_bundles: Query>, ) { debug!("queue_gaussians"); - let warmup = views.iter().any(|(_, camera, _, _)| camera.warmup); + let warmup = views.iter().any(|(_, camera, _, _, _, _)| camera.warmup); if warmup { debug!("skipping gaussian cloud render during warmup"); return; @@ -375,7 +414,7 @@ fn queue_gaussians( .read() .id::>(); - for (view, _, visible_entities, msaa) in &mut views { + for (view, _, visible_entities, msaa, prepass_textures, has_oit) in &mut views { debug!("queue gaussians view"); let Some(transparent_phase) = transparent_render_phases.get_mut(&view.retained_view_entity) else { @@ -404,7 +443,12 @@ fn queue_gaussians( return; } - let msaa = msaa.cloned().unwrap_or_default(); + let msaa = msaa.map(|m| m.clone()).unwrap_or_default(); + let mut mesh_view_layout_key = + MeshPipelineViewLayoutKey::from(msaa) | MeshPipelineViewLayoutKey::from(prepass_textures); + if has_oit { + mesh_view_layout_key |= MeshPipelineViewLayoutKey::OIT_ENABLED; + } let key = CloudPipelineKey { aabb: settings.aabb, @@ -416,6 +460,7 @@ fn queue_gaussians( rasterize_mode: settings.rasterize_mode, sample_count: msaa.samples(), hdr: view.hdr, + mesh_view_layout_key_bits: mesh_view_layout_key.bits(), }; let pipeline = pipelines.specialize(&pipeline_cache, &custom_pipeline, key); @@ -425,8 +470,7 @@ fn queue_gaussians( let aabb_size = aabb.max() - aabb.min(); let center = *transform * GlobalTransform::from( - Transform::from_translation(aabb_center.into()) - .with_scale(aabb_size.into()), + Transform::from_translation(aabb_center.into()).with_scale(aabb_size.into()), ); let distance = rangefinder.distance_translation(¢er.translation()); @@ -453,6 +497,15 @@ pub struct CloudPipeline { pub view_layout: BindGroupLayout, pub compute_view_layout: BindGroupLayout, pub sorted_layout: BindGroupLayout, + pub pbr_layout: BindGroupLayout, + pub fallback_normals: Buffer, + pub fallback_materials: Buffer, + pub fallback_bind_group: BindGroup, + pub gaussian_material_layout: BindGroupLayout, + pub _fallback_material_override_uniform: Buffer, + pub fallback_material_override_bind_group: BindGroup, + pub _fallback_material_override_sampler: Sampler, + mesh_view_layouts: MeshPipelineViewLayouts, phantom: std::marker::PhantomData, } @@ -606,6 +659,139 @@ where #[cfg(feature = "buffer_texture")] let sorted_layout = texture::get_sorted_bind_group_layout(render_device); + let pbr_layout = render_device.create_bind_group_layout( + Some("gaussian_pbr_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::VERTEX_FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: Some(NormalData::min_size()), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::VERTEX_FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: Some(PbrMaterialData::min_size()), + }, + count: None, + }, + ], + ); + + let gaussian_material_layout = render_device.create_bind_group_layout( + Some("gaussian_material_override_layout"), + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Texture { + sample_type: TextureSampleType::Float { filterable: true }, + view_dimension: TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Sampler(SamplerBindingType::Filtering), + count: None, + }, + ], + ); + + let fallback_normal = NormalData::default(); + let fallback_normals = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("gaussian_pbr_fallback_normals"), + contents: bytes_of(&fallback_normal), + usage: BufferUsages::STORAGE, + }); + + let fallback_material = PbrMaterialData::default(); + let fallback_materials = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("gaussian_pbr_fallback_materials"), + contents: bytes_of(&fallback_material), + usage: BufferUsages::STORAGE, + }); + + let fallback_bind_group = render_device.create_bind_group( + "gaussian_pbr_fallback_bind_group", + &pbr_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &fallback_normals, + offset: 0, + size: None, + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::Buffer(BufferBinding { + buffer: &fallback_materials, + offset: 0, + size: None, + }), + }, + ], + ); + + let fallback_image = render_world.resource::(); + + let fallback_override = GaussianMaterialOverrideUniform { + base_color_factor: [1.0, 1.0, 1.0, 1.0], + bounds_min: [0.0, 0.0, 0.0, 0.0], + bounds_size: [1.0, 1.0, 1.0, 0.0], + flags: [0, 0, 0, 0], + }; + + let fallback_material_override_uniform = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("gaussian_material_override_uniform"), + contents: bytes_of(&fallback_override), + usage: BufferUsages::UNIFORM, + }); + + let fallback_material_override_sampler = fallback_image.d2.sampler.clone(); + + let fallback_material_override_bind_group = render_device.create_bind_group( + "gaussian_material_override_fallback", + &gaussian_material_layout, + &[ + BindGroupEntry { + binding: 0, + resource: fallback_material_override_uniform.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::TextureView(&fallback_image.d2.texture_view), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::Sampler(&fallback_material_override_sampler), + }, + ], + ); + + let mesh_view_layouts = render_world.resource::().clone(); + debug!("created cloud pipeline"); Self { @@ -615,6 +801,15 @@ where compute_view_layout, shader: GAUSSIAN_SHADER_HANDLE, sorted_layout, + pbr_layout, + fallback_normals, + fallback_materials, + fallback_bind_group, + gaussian_material_layout, + _fallback_material_override_uniform: fallback_material_override_uniform, + fallback_material_override_bind_group, + _fallback_material_override_sampler: fallback_material_override_sampler, + mesh_view_layouts, phantom: std::marker::PhantomData, } } @@ -815,6 +1010,7 @@ pub struct CloudPipelineKey { pub rasterize_mode: RasterizeMode, pub sample_count: u32, pub hdr: bool, + pub mesh_view_layout_key_bits: u32, } impl SpecializedRenderPipeline for CloudPipeline { @@ -831,13 +1027,19 @@ impl SpecializedRenderPipeline for CloudPipeline { debug!("specializing cloud pipeline"); + let view_layout_key = + MeshPipelineViewLayoutKey::from_bits_truncate(key.mesh_view_layout_key_bits); + let view_layout = &self.mesh_view_layouts[view_layout_key.bits() as usize].main_layout; + RenderPipelineDescriptor { label: Some("gaussian cloud render pipeline".into()), layout: vec![ - self.view_layout.clone(), + view_layout.clone(), self.gaussian_uniform_layout.clone(), self.gaussian_cloud_layout.clone(), self.sorted_layout.clone(), + self.pbr_layout.clone(), + self.gaussian_material_layout.clone(), ], vertex: VertexState { shader: self.shader.clone(), @@ -893,8 +1095,7 @@ impl SpecializedRenderPipeline for CloudPipeline { type DrawGaussians = ( SetItemPipeline, - // SetViewBindGroup<0>, - SetPreviousViewBindGroup<0>, + SetMeshViewBindGroup<0>, SetGaussianUniformBindGroup<1>, DrawGaussianInstanced, ); @@ -1000,6 +1201,21 @@ pub struct SortBindGroup { pub sorted_bind_group: BindGroup, } +#[derive(Component)] +pub struct GaussianPbrBindGroup { + pub bind_group: BindGroup, +} + +#[derive(Component)] +pub struct GaussianMaterialOverrideBindGroup { + pub bind_group: BindGroup, + pub uniform: Buffer, + pub asset_id: Option>, + pub revision: u64, + pub bounds_min: Vec3, + pub bounds_size: Vec3, +} + #[allow(clippy::too_many_arguments)] fn queue_gaussian_bind_group( mut commands: Commands, @@ -1118,6 +1334,157 @@ fn queue_gaussian_bind_group( } } +fn queue_gaussian_pbr_bind_group_3d( + mut commands: Commands, + render_device: Res, + gaussian_cloud_pipeline: Res>, + normals_res: Res, + materials_res: Res, + material_assets: Res, + gaussian_clouds: Query<( + Entity, + &PlanarGaussian3dHandle, + &CloudUniform, + Option<&GaussianPbrBindGroup>, + Option<&GaussianMaterialOverrideBindGroup>, + Option<&GaussianMaterialHandle>, + )>, +) { + let pipeline_changed = gaussian_cloud_pipeline.is_changed(); + + for (entity, planar_handle, cloud_uniform, existing_pbr, existing_override, material_handle) in + gaussian_clouds.iter() + { + if existing_pbr.is_some() && !pipeline_changed { + continue; + } + + let asset_id = planar_handle.handle().id(); + let Some(normals_gpu) = normals_res.asset_map.get(&asset_id) else { + continue; + }; + let Some(materials_gpu) = materials_res.asset_map.get(&asset_id) else { + continue; + }; + + let bind_group = render_device.create_bind_group( + "gaussian_pbr_bind_group", + &gaussian_cloud_pipeline.pbr_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &normals_gpu.normals, + offset: 0, + size: None, + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::Buffer(BufferBinding { + buffer: &materials_gpu.materials, + offset: 0, + size: None, + }), + }, + ], + ); + + commands + .entity(entity) + .insert(GaussianPbrBindGroup { bind_group }); + + // Material override handling + let Some(material_asset_id) = material_handle + .filter(|h| h.is_strong()) + .map(|h| h.id()) + else { + if existing_override.is_some() { + commands + .entity(entity) + .remove::(); + } + continue; + }; + + let Some(cached_material) = material_assets.map.get(&material_asset_id) else { + if existing_override.is_some() { + commands + .entity(entity) + .remove::(); + } + continue; + }; + + let gpu_material = &cached_material.material; + + let bounds = gpu_material + .bounds_override + .unwrap_or(GaussianMaterialBounds { + min: cloud_uniform.min.truncate(), + max: cloud_uniform.max.truncate(), + }); + + let bounds_min = bounds.min; + let bounds_size = (bounds.max - bounds.min).max(Vec3::splat(1e-3)); + + let needs_override_update = pipeline_changed + || existing_override + .map(|override_bg| { + override_bg.asset_id != Some(material_asset_id) + || override_bg.bounds_min != bounds_min + || override_bg.bounds_size != bounds_size + || override_bg.revision != cached_material.revision + }) + .unwrap_or(true); + + if !needs_override_update { + continue; + } + + let override_uniform = GaussianMaterialOverrideUniform { + base_color_factor: gpu_material.base_color.into(), + bounds_min: bounds_min.extend(0.0).to_array(), + bounds_size: bounds_size.extend(0.0).to_array(), + flags: [gpu_material.use_texture, gpu_material.projection_axis, 0, 0], + }; + + let uniform_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: Some("gaussian_material_override_uniform"), + contents: bytes_of(&override_uniform), + usage: BufferUsages::UNIFORM, + }); + + let override_bind_group = render_device.create_bind_group( + "gaussian_material_override_bind_group", + &gaussian_cloud_pipeline.gaussian_material_layout, + &[ + BindGroupEntry { + binding: 0, + resource: uniform_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::TextureView(&gpu_material.texture_view), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::Sampler(&gpu_material.sampler), + }, + ], + ); + + commands.entity(entity).insert(GaussianMaterialOverrideBindGroup { + bind_group: override_bind_group, + uniform: uniform_buffer, + asset_id: Some(material_asset_id), + revision: cached_material.revision, + bounds_min, + bounds_size, + }); + } +} + #[derive(Component)] pub struct GaussianViewBindGroup { pub value: BindGroup, @@ -1284,74 +1651,6 @@ pub fn queue_gaussian_compute_view_bind_groups( } } -pub struct SetViewBindGroup; -impl RenderCommand

for SetViewBindGroup { - type Param = (); - type ViewQuery = (Read, Read); - type ItemQuery = (); - - #[inline] - fn render<'w>( - _: &P, - (gaussian_view_bind_group, view_uniform): ROQueryItem<'w, 'w, Self::ViewQuery>, - _entity: Option<()>, - _: SystemParamItem<'w, '_, Self::Param>, - pass: &mut TrackedRenderPass<'w>, - ) -> RenderCommandResult { - pass.set_bind_group(I, &gaussian_view_bind_group.value, &[view_uniform.offset]); - - debug!("set view bind group"); - - RenderCommandResult::Success - } -} - -pub struct SetPreviousViewBindGroup; -impl RenderCommand

for SetPreviousViewBindGroup { - type Param = SRes; - type ViewQuery = ( - Read, - Option>, - Option>, - ); - type ItemQuery = (); - - #[inline] - fn render<'w>( - _: &P, - (view_uniform_offset, has_motion_vector_prepass, previous_view_uniform_offset): ROQueryItem< - 'w, - 'w, - Self::ViewQuery, - >, - _entity: Option<()>, - prepass_view_bind_group: SystemParamItem<'w, '_, Self::Param>, - pass: &mut TrackedRenderPass<'w>, - ) -> RenderCommandResult { - let prepass_view_bind_group = prepass_view_bind_group.into_inner(); - match previous_view_uniform_offset { - Some(previous_view_uniform_offset) if has_motion_vector_prepass.unwrap_or_default() => { - pass.set_bind_group( - I, - prepass_view_bind_group.motion_vectors.as_ref().unwrap(), - &[ - view_uniform_offset.offset, - previous_view_uniform_offset.offset, - ], - ); - } - _ => pass.set_bind_group( - I, - prepass_view_bind_group.motion_vectors.as_ref().unwrap(), - &[view_uniform_offset.offset, 0], - ), - } - - debug!("set previous view bind group"); - - RenderCommandResult::Success - } -} pub struct SetGaussianUniformBindGroup; impl RenderCommand

for SetGaussianUniformBindGroup { @@ -1405,12 +1704,17 @@ impl RenderCommand

for DrawGaussianInstanced where R::GpuPlanarType: GpuPlanarStorage, { - type Param = SRes>; + type Param = ( + SRes>, + SRes>, + ); type ViewQuery = Read; type ItemQuery = ( Read, Read>, Read, + Option>, + Option>, ); #[inline] @@ -1421,16 +1725,27 @@ where &'w R::PlanarTypeHandle, &'w PlanarStorageBindGroup, &'w SortBindGroup, + Option<&'w GaussianPbrBindGroup>, + Option<&'w GaussianMaterialOverrideBindGroup>, )>, - gaussian_clouds: SystemParamItem<'w, '_, Self::Param>, + gaussian_resources: SystemParamItem<'w, '_, Self::Param>, pass: &mut TrackedRenderPass<'w>, ) -> RenderCommandResult { debug!("render call"); - let (handle, planar_bind_groups, sort_bind_groups) = + let gaussian_assets: &RenderAssets = gaussian_resources.0.into_inner(); + let pipeline: &CloudPipeline = gaussian_resources.1.into_inner(); + + let ( + handle, + planar_bind_groups, + sort_bind_groups, + maybe_pbr_bind_group, + maybe_material_override, + ) = entity.expect("gaussian cloud entity not found"); - let gpu_gaussian_cloud = match gaussian_clouds.into_inner().get(handle.handle()) { + let gpu_gaussian_cloud = match gaussian_assets.get(handle.handle()) { Some(gpu_gaussian_cloud) => gpu_gaussian_cloud, None => { debug!("gpu cloud not found"); @@ -1451,6 +1766,16 @@ where * gpu_gaussian_cloud.len() as u32], ); + let pbr_bind_group = maybe_pbr_bind_group + .map(|bg| &bg.bind_group) + .unwrap_or(&pipeline.fallback_bind_group); + pass.set_bind_group(4, pbr_bind_group, &[]); + + let material_override_bind_group = maybe_material_override + .map(|bg| &bg.bind_group) + .unwrap_or(&pipeline.fallback_material_override_bind_group); + pass.set_bind_group(5, material_override_bind_group, &[]); + #[cfg(feature = "webgl2")] pass.draw(0..4, 0..gpu_gaussian_cloud.count as u32); diff --git a/src/render/transform.wgsl b/src/render/transform.wgsl index 72f455c0..85b24f74 100644 --- a/src/render/transform.wgsl +++ b/src/render/transform.wgsl @@ -1,6 +1,6 @@ #define_import_path bevy_gaussian_splatting::transform -#import bevy_gaussian_splatting::bindings::view +#import bevy_pbr::mesh_view_bindings::view fn world_to_clip(world_pos: vec3) -> vec4 { let homogenous_pos = view.unjittered_clip_from_world * vec4(world_pos, 1.0); diff --git a/src/sort/radix.wgsl b/src/sort/radix.wgsl index d15d2a29..fb6ffa49 100644 --- a/src/sort/radix.wgsl +++ b/src/sort/radix.wgsl @@ -1,5 +1,4 @@ #import bevy_gaussian_splatting::bindings::{ - view, globals, gaussian_uniforms, sorting_pass_index, @@ -12,6 +11,7 @@ DrawIndirect, Entry, } +#import bevy_pbr::mesh_view_bindings as view_bindings; #import bevy_gaussian_splatting::transform::{ world_to_clip, in_frustum, @@ -90,7 +90,7 @@ fn radix_sort_a( let position = vec4(get_position(entry_index), 1.0); let transformed_position = (gaussian_uniforms.transform * position).xyz; let clip_space_pos = world_to_clip(transformed_position); - let diff = transformed_position - view.world_position; + let diff = transformed_position - view_bindings::view.world_position; let dist2 = dot(diff, diff); let dist_bits = bitcast(dist2); let key_distance = 0xFFFFFFFFu - dist_bits; diff --git a/tools/generate_ggx_lut.rs b/tools/generate_ggx_lut.rs new file mode 100644 index 00000000..e573184a --- /dev/null +++ b/tools/generate_ggx_lut.rs @@ -0,0 +1,205 @@ +use std::f32::consts::PI; +use std::path::Path; + +const LUT_SIZE: u32 = 128; +const SAMPLE_COUNT: u32 = 256; + +fn main() { + println!("Generating GGX BRDF integration LUT..."); + println!("Resolution: {}x{}", LUT_SIZE, LUT_SIZE); + println!("Samples per pixel: {}", SAMPLE_COUNT); + + let mut lut_data = vec![0u8; (LUT_SIZE * LUT_SIZE * 4) as usize]; + + for y in 0..LUT_SIZE { + for x in 0..LUT_SIZE { + let roughness = (x as f32 + 0.5) / LUT_SIZE as f32; + let ndotv = (y as f32 + 0.5) / LUT_SIZE as f32; + + let (scale, bias) = integrate_ggx_brdf(roughness, ndotv); + + let idx = ((y * LUT_SIZE + x) * 4) as usize; + + let scale_u16 = (scale * 65535.0).clamp(0.0, 65535.0) as u16; + let bias_u16 = (bias * 65535.0).clamp(0.0, 65535.0) as u16; + + lut_data[idx] = (scale_u16 & 0xFF) as u8; + lut_data[idx + 1] = ((scale_u16 >> 8) & 0xFF) as u8; + lut_data[idx + 2] = (bias_u16 & 0xFF) as u8; + lut_data[idx + 3] = ((bias_u16 >> 8) & 0xFF) as u8; + } + + if y % 16 == 0 { + println!("Progress: {}/{}", y, LUT_SIZE); + } + } + + let output_path = Path::new("assets/textures/ggx_energy_lut.png"); + if let Some(parent) = output_path.parent() { + std::fs::create_dir_all(parent).expect("Failed to create output directory"); + } + + save_lut_as_png(&lut_data, output_path); + + println!("LUT generated successfully: {:?}", output_path); +} + +fn integrate_ggx_brdf(perceptual_roughness: f32, ndotv: f32) -> (f32, f32) { + let roughness = perceptual_roughness * perceptual_roughness; + + let v = Vec3::new( + (1.0 - ndotv * ndotv).sqrt().max(0.0), + 0.0, + ndotv, + ); + + let mut scale = 0.0; + let mut bias = 0.0; + + for i in 0..SAMPLE_COUNT { + let xi = hammersley(i, SAMPLE_COUNT); + let h = importance_sample_ggx(xi, roughness); + let l = 2.0 * h.dot(&v) * h - v; + + let ndotl = l.z.max(0.0); + let ndoth = h.z.max(0.0); + let vdoth = v.dot(&h).max(0.0); + + if ndotl > 0.0 { + let g = geometry_smith(ndotv, ndotl, roughness); + let g_vis = g * vdoth / (ndoth * ndotv).max(1e-6); + let fc = (1.0 - vdoth).powf(5.0); + + scale += (1.0 - fc) * g_vis; + bias += fc * g_vis; + } + } + + scale /= SAMPLE_COUNT as f32; + bias /= SAMPLE_COUNT as f32; + + (scale, bias) +} + +fn hammersley(i: u32, n: u32) -> Vec2 { + Vec2::new( + i as f32 / n as f32, + radical_inverse_vdc(i), + ) +} + +fn radical_inverse_vdc(mut bits: u32) -> f32 { + bits = (bits << 16) | (bits >> 16); + bits = ((bits & 0x55555555) << 1) | ((bits & 0xAAAAAAAA) >> 1); + bits = ((bits & 0x33333333) << 2) | ((bits & 0xCCCCCCCC) >> 2); + bits = ((bits & 0x0F0F0F0F) << 4) | ((bits & 0xF0F0F0F0) >> 4); + bits = ((bits & 0x00FF00FF) << 8) | ((bits & 0xFF00FF00) >> 8); + bits as f32 * 2.3283064365386963e-10 +} + +fn importance_sample_ggx(xi: Vec2, roughness: f32) -> Vec3 { + let a = roughness * roughness; + + let phi = 2.0 * PI * xi.x; + let cos_theta = ((1.0 - xi.y) / (1.0 + (a * a - 1.0) * xi.y)).sqrt(); + let sin_theta = (1.0 - cos_theta * cos_theta).sqrt(); + + Vec3::new( + phi.cos() * sin_theta, + phi.sin() * sin_theta, + cos_theta, + ) +} + +fn geometry_schlick_ggx(ndotv: f32, roughness: f32) -> f32 { + let a = roughness; + let k = (a * a) / 2.0; + + let nom = ndotv; + let denom = ndotv * (1.0 - k) + k; + + nom / denom.max(1e-6) +} + +fn geometry_smith(ndotv: f32, ndotl: f32, roughness: f32) -> f32 { + let ggx2 = geometry_schlick_ggx(ndotv, roughness); + let ggx1 = geometry_schlick_ggx(ndotl, roughness); + + ggx1 * ggx2 +} + +fn save_lut_as_png(data: &[u8], path: &Path) { + use std::fs::File; + use std::io::BufWriter; + + let file = File::create(path).expect("Failed to create output file"); + let w = BufWriter::new(file); + + let mut encoder = png::Encoder::new(w, LUT_SIZE, LUT_SIZE); + encoder.set_color(png::ColorType::Rgba); + encoder.set_depth(png::BitDepth::Eight); + + let mut writer = encoder.write_header().expect("Failed to write PNG header"); + writer.write_image_data(data).expect("Failed to write PNG data"); +} + +#[derive(Debug, Clone, Copy)] +struct Vec2 { + x: f32, + y: f32, +} + +impl Vec2 { + fn new(x: f32, y: f32) -> Self { + Self { x, y } + } +} + +#[derive(Debug, Clone, Copy)] +struct Vec3 { + x: f32, + y: f32, + z: f32, +} + +impl Vec3 { + fn new(x: f32, y: f32, z: f32) -> Self { + Self { x, y, z } + } + + fn dot(&self, other: &Self) -> f32 { + self.x * other.x + self.y * other.y + self.z * other.z + } +} + +impl std::ops::Sub for Vec3 { + type Output = Self; + + fn sub(self, other: Self) -> Self { + Self { + x: self.x - other.x, + y: self.y - other.y, + z: self.z - other.z, + } + } +} + +impl std::ops::Mul for Vec3 { + type Output = Self; + + fn mul(self, scalar: f32) -> Self { + Self { + x: self.x * scalar, + y: self.y * scalar, + z: self.z * scalar, + } + } +} + +impl std::ops::Mul for f32 { + type Output = Vec3; + + fn mul(self, vec: Vec3) -> Vec3 { + vec * self + } +}