docs/bevy/examples/3d/reflection_probes.rs 13.7 K · 429 lines · raw · history

1 //! This example shows how to place reflection probes in the scene.
2 //!
3 //! Press Space to cycle through the reflection modes:
4 //!
5 //! 1. A pre-generated [`EnvironmentMapLight`] acting as a reflection probe, with both the skybox and cubes
6 //! 2. A runtime-generated [`GeneratedEnvironmentMapLight`] acting as a reflection probe with just the skybox
7 //! 3. A pre-generated [`EnvironmentMapLight`] with just the skybox
8 //!
9 //! Press Enter to pause or resume rotation.
10 //!
11 //! Reflection probes don't work on WebGL 2 or WebGPU.
12
13 use bevy::{
14 camera::{Exposure, Hdr},
15 core_pipeline::tonemapping::Tonemapping,
16 light::{ParallaxCorrection, Skybox},
17 pbr::generate::generate_environment_map_light,
18 prelude::*,
19 render::render_resource::TextureUsages,
20 };
21
22 use std::{
23 f32::consts::PI,
24 fmt::{Display, Formatter, Result as FmtResult},
25 };
26
27 static STOP_ROTATION_HELP_TEXT: &str = "Press Enter to stop rotation";
28 static START_ROTATION_HELP_TEXT: &str = "Press Enter to start rotation";
29
30 static REFLECTION_MODE_HELP_TEXT: &str = "Press Space to switch reflection mode";
31
32 const ENV_MAP_INTENSITY: f32 = 5000.0;
33
34 // The mode the application is in.
35 #[derive(Resource)]
36 struct AppStatus {
37 // Which environment maps the user has requested to display.
38 reflection_mode: ReflectionMode,
39 // Whether the user has requested the scene to rotate.
40 rotating: bool,
41 // The current roughness of the central sphere
42 sphere_roughness: f32,
43 }
44
45 // Which environment maps the user has requested to display.
46 #[derive(Clone, Copy, PartialEq)]
47 enum ReflectionMode {
48 // Only a world environment map is shown.
49 EnvironmentMap = 0,
50 // Both a world environment map and a reflection probe are present. The
51 // reflection probe is shown in the sphere.
52 ReflectionProbe = 1,
53 // A generated environment map is shown.
54 GeneratedEnvironmentMap = 2,
55 }
56
57 // The various reflection maps.
58 #[derive(Resource)]
59 struct Cubemaps {
60 // The blurry diffuse cubemap that reflects the world, but not the cubes.
61 diffuse_environment_map: Handle<Image>,
62
63 // The specular cubemap mip chain that reflects the world, but not the cubes.
64 specular_environment_map: Handle<Image>,
65
66 // The specular cubemap mip chain that reflects both the world and the cubes.
67 specular_reflection_probe: Handle<Image>,
68 }
69
70 fn main() {
71 // Create the app.
72 App::new()
73 .add_plugins(DefaultPlugins)
74 .init_resource::<AppStatus>()
75 .init_resource::<Cubemaps>()
76 .add_systems(Startup, setup)
77 .add_systems(PreUpdate, add_environment_map_to_camera)
78 .add_systems(
79 Update,
80 change_reflection_type.before(generate_environment_map_light),
81 )
82 .add_systems(Update, toggle_rotation)
83 .add_systems(Update, change_sphere_roughness)
84 .add_systems(
85 Update,
86 rotate_camera
87 .after(toggle_rotation)
88 .after(change_reflection_type),
89 )
90 .add_systems(Update, update_text.after(rotate_camera))
91 .add_systems(Update, setup_environment_map_usage)
92 .run();
93 }
94
95 // Spawns all the scene objects.
96 fn setup(
97 mut commands: Commands,
98 mut meshes: ResMut<Assets<Mesh>>,
99 mut materials: ResMut<Assets<StandardMaterial>>,
100 asset_server: Res<AssetServer>,
101 app_status: Res<AppStatus>,
102 cubemaps: Res<Cubemaps>,
103 ) {
104 spawn_camera(&mut commands);
105 spawn_sphere(&mut commands, &mut meshes, &mut materials, &app_status);
106 spawn_reflection_probe(&mut commands, &cubemaps);
107 spawn_scene(&mut commands, &asset_server);
108 spawn_text(&mut commands, &app_status);
109 }
110
111 // Spawns the cubes, light, and camera.
112 fn spawn_scene(commands: &mut Commands, asset_server: &AssetServer) {
113 commands.spawn((
114 WorldAssetRoot(
115 asset_server.load(GltfAssetLabel::Scene(0).from_asset("models/cubes/Cubes.glb")),
116 ),
117 CubesScene,
118 ));
119 }
120
121 // Spawns the camera.
122 fn spawn_camera(commands: &mut Commands) {
123 commands.spawn((
124 Camera3d::default(),
125 Hdr,
126 Exposure { ev100: 11.0 },
127 Tonemapping::AcesFitted,
128 Transform::from_xyz(-3.883, 0.325, 2.781).looking_at(Vec3::ZERO, Vec3::Y),
129 ));
130 }
131
132 // Creates the sphere mesh and spawns it.
133 fn spawn_sphere(
134 commands: &mut Commands,
135 meshes: &mut Assets<Mesh>,
136 materials: &mut Assets<StandardMaterial>,
137 app_status: &AppStatus,
138 ) {
139 // Create a sphere mesh.
140 let sphere_mesh = meshes.add(Sphere::new(1.0).mesh().ico(7).unwrap());
141
142 // Create a sphere.
143 commands.spawn((
144 Mesh3d(sphere_mesh.clone()),
145 MeshMaterial3d(materials.add(StandardMaterial {
146 base_color: Srgba::hex("#ffffff").unwrap().into(),
147 metallic: 1.0,
148 perceptual_roughness: app_status.sphere_roughness,
149 ..StandardMaterial::default()
150 })),
151 SphereMaterial,
152 ));
153 }
154
155 // Spawns the reflection probe.
156 fn spawn_reflection_probe(commands: &mut Commands, cubemaps: &Cubemaps) {
157 commands.spawn((
158 LightProbe::default(),
159 EnvironmentMapLight {
160 diffuse_map: cubemaps.diffuse_environment_map.clone(),
161 specular_map: cubemaps.specular_reflection_probe.clone(),
162 intensity: ENV_MAP_INTENSITY,
163 ..default()
164 },
165 // 2.0 because the sphere's radius is 1.0 and we want to fully enclose it.
166 Transform::from_scale(Vec3::splat(2.0)),
167 // Disable parallax correction because the reflected scene is quite
168 // distant.
169 ParallaxCorrection::None,
170 ));
171 }
172
173 // Spawns the help text.
174 fn spawn_text(commands: &mut Commands, app_status: &AppStatus) {
175 // Create the text.
176 commands.spawn((
177 app_status.create_text(),
178 Node {
179 position_type: PositionType::Absolute,
180 bottom: px(12),
181 left: px(12),
182 ..default()
183 },
184 ));
185 }
186
187 // Adds a world environment map to the camera. This separate system is needed because the camera is
188 // managed by the scene spawner, as it's part of the glTF file with the cubes, so we have to add
189 // the environment map after the fact.
190 fn add_environment_map_to_camera(
191 mut commands: Commands,
192 query: Query<Entity, Added<Camera3d>>,
193 cubemaps: Res<Cubemaps>,
194 ) {
195 for camera_entity in query.iter() {
196 commands
197 .entity(camera_entity)
198 .insert(create_camera_environment_map_light(&cubemaps))
199 .insert(Skybox {
200 image: Some(cubemaps.specular_environment_map.clone()),
201 brightness: ENV_MAP_INTENSITY,
202 ..default()
203 });
204 }
205 }
206
207 // A system that handles switching between different reflection modes.
208 fn change_reflection_type(
209 mut commands: Commands,
210 light_probe_query: Query<Entity, With<LightProbe>>,
211 cubes_scene_query: Query<Entity, With<CubesScene>>,
212 camera_query: Query<Entity, With<Camera3d>>,
213 keyboard: Res<ButtonInput<KeyCode>>,
214 mut app_status: ResMut<AppStatus>,
215 cubemaps: Res<Cubemaps>,
216 asset_server: Res<AssetServer>,
217 ) {
218 // Only do anything if space was pressed.
219 if !keyboard.just_pressed(KeyCode::Space) {
220 return;
221 }
222
223 // Advance to the next reflection mode.
224 app_status.reflection_mode =
225 ReflectionMode::try_from((app_status.reflection_mode as u32 + 1) % 3).unwrap();
226
227 // Remove light probes
228 for light_probe in light_probe_query.iter() {
229 commands.entity(light_probe).despawn();
230 }
231 // Remove existing cube scenes
232 for scene_entity in cubes_scene_query.iter() {
233 commands.entity(scene_entity).despawn();
234 }
235 match app_status.reflection_mode {
236 ReflectionMode::EnvironmentMap | ReflectionMode::GeneratedEnvironmentMap => {}
237 ReflectionMode::ReflectionProbe => {
238 spawn_reflection_probe(&mut commands, &cubemaps);
239 spawn_scene(&mut commands, &asset_server);
240 }
241 }
242
243 // Update the environment-map components on the camera entity/entities
244 for camera in camera_query.iter() {
245 // Remove any existing environment-map components
246 commands
247 .entity(camera)
248 .remove::<(EnvironmentMapLight, GeneratedEnvironmentMapLight)>();
249
250 match app_status.reflection_mode {
251 // A baked or reflection-probe environment map
252 ReflectionMode::EnvironmentMap | ReflectionMode::ReflectionProbe => {
253 commands
254 .entity(camera)
255 .insert(create_camera_environment_map_light(&cubemaps));
256 }
257
258 // GPU-filtered environment map generated at runtime
259 ReflectionMode::GeneratedEnvironmentMap => {
260 commands
261 .entity(camera)
262 .insert(GeneratedEnvironmentMapLight {
263 environment_map: cubemaps.specular_environment_map.clone(),
264 intensity: ENV_MAP_INTENSITY,
265 ..default()
266 });
267 }
268 }
269 }
270 }
271
272 // A system that handles enabling and disabling rotation.
273 fn toggle_rotation(keyboard: Res<ButtonInput<KeyCode>>, mut app_status: ResMut<AppStatus>) {
274 if keyboard.just_pressed(KeyCode::Enter) {
275 app_status.rotating = !app_status.rotating;
276 }
277 }
278
279 // A system that updates the help text.
280 fn update_text(mut text_query: Query<&mut Text>, app_status: Res<AppStatus>) {
281 for mut text in text_query.iter_mut() {
282 *text = app_status.create_text();
283 }
284 }
285
286 impl TryFrom<u32> for ReflectionMode {
287 type Error = ();
288
289 fn try_from(value: u32) -> Result<Self, Self::Error> {
290 match value {
291 0 => Ok(ReflectionMode::EnvironmentMap),
292 1 => Ok(ReflectionMode::ReflectionProbe),
293 2 => Ok(ReflectionMode::GeneratedEnvironmentMap),
294 _ => Err(()),
295 }
296 }
297 }
298
299 impl Display for ReflectionMode {
300 fn fmt(&self, formatter: &mut Formatter<'_>) -> FmtResult {
301 let text = match *self {
302 ReflectionMode::EnvironmentMap => "Environment map",
303 ReflectionMode::ReflectionProbe => "Reflection probe",
304 ReflectionMode::GeneratedEnvironmentMap => "Generated environment map",
305 };
306 formatter.write_str(text)
307 }
308 }
309
310 impl AppStatus {
311 // Constructs the help text at the bottom of the screen based on the
312 // application status.
313 fn create_text(&self) -> Text {
314 let rotation_help_text = if self.rotating {
315 STOP_ROTATION_HELP_TEXT
316 } else {
317 START_ROTATION_HELP_TEXT
318 };
319
320 format!(
321 "{}\n{}\nRoughness: {:.2}\n{}\nUp/Down arrows to change roughness",
322 self.reflection_mode,
323 rotation_help_text,
324 self.sphere_roughness,
325 REFLECTION_MODE_HELP_TEXT
326 )
327 .into()
328 }
329 }
330
331 // Creates the world environment map light, used as a fallback if no reflection
332 // probe is applicable to a mesh.
333 fn create_camera_environment_map_light(cubemaps: &Cubemaps) -> EnvironmentMapLight {
334 EnvironmentMapLight {
335 diffuse_map: cubemaps.diffuse_environment_map.clone(),
336 specular_map: cubemaps.specular_environment_map.clone(),
337 intensity: ENV_MAP_INTENSITY,
338 ..default()
339 }
340 }
341
342 // Rotates the camera a bit every frame.
343 fn rotate_camera(
344 time: Res<Time>,
345 mut camera_query: Query<&mut Transform, With<Camera3d>>,
346 app_status: Res<AppStatus>,
347 ) {
348 if !app_status.rotating {
349 return;
350 }
351
352 for mut transform in camera_query.iter_mut() {
353 transform.translation = Vec2::from_angle(time.delta_secs() * PI / 5.0)
354 .rotate(transform.translation.xz())
355 .extend(transform.translation.y)
356 .xzy();
357 transform.look_at(Vec3::ZERO, Vec3::Y);
358 }
359 }
360
361 // Loads the cubemaps from the assets directory.
362 impl FromWorld for Cubemaps {
363 fn from_world(world: &mut World) -> Self {
364 Cubemaps {
365 diffuse_environment_map: world
366 .load_asset("environment_maps/pisa_diffuse_rgb9e5_zstd.ktx2"),
367 specular_environment_map: world
368 .load_asset("environment_maps/pisa_specular_rgb9e5_zstd.ktx2"),
369 specular_reflection_probe: world
370 .load_asset("environment_maps/cubes_reflection_probe_specular_rgb9e5_zstd.ktx2"),
371 }
372 }
373 }
374
375 fn setup_environment_map_usage(cubemaps: Res<Cubemaps>, mut images: ResMut<Assets<Image>>) {
376 if let Some(mut image) = images.get_mut(&cubemaps.specular_environment_map)
377 && !image
378 .texture_descriptor
379 .usage
380 .contains(TextureUsages::COPY_SRC)
381 {
382 image.texture_descriptor.usage |= TextureUsages::COPY_SRC;
383 }
384 }
385
386 impl Default for AppStatus {
387 fn default() -> Self {
388 Self {
389 reflection_mode: ReflectionMode::ReflectionProbe,
390 rotating: false,
391 sphere_roughness: 0.2,
392 }
393 }
394 }
395
396 #[derive(Component)]
397 struct SphereMaterial;
398
399 #[derive(Component)]
400 struct CubesScene;
401
402 // A system that changes the sphere's roughness with up/down arrow keys
403 fn change_sphere_roughness(
404 keyboard: Res<ButtonInput<KeyCode>>,
405 mut app_status: ResMut<AppStatus>,
406 mut materials: ResMut<Assets<StandardMaterial>>,
407 sphere_query: Query<&MeshMaterial3d<StandardMaterial>, With<SphereMaterial>>,
408 ) {
409 let roughness_delta = if keyboard.pressed(KeyCode::ArrowUp) {
410 0.01 // Decrease roughness
411 } else if keyboard.pressed(KeyCode::ArrowDown) {
412 -0.01 // Increase roughness
413 } else {
414 0.0 // No change
415 };
416
417 if roughness_delta != 0.0 {
418 // Update the app status
419 app_status.sphere_roughness =
420 (app_status.sphere_roughness + roughness_delta).clamp(0.0, 1.0);
421
422 // Update the sphere material
423 for material_handle in sphere_query.iter() {
424 if let Some(mut material) = materials.get_mut(&material_handle.0) {
425 material.perceptual_roughness = app_status.sphere_roughness;
426 }
427 }
428 }
429 }