Bevy v0.10 draw Axis-Aligned Bounding-Boxes (AABB) around Scenes after they've been loaded

scene_aabb.rs

//! Load a GLB Scene, and draw an Axis-Aligned Bounding-Box around the entire
//! scene
//!
//! Works in Bevy v0.10
#![allow(clippy::type_complexity)]

use bevy::{
    core_pipeline::tonemapping::Tonemapping, prelude::*, render::primitives::Aabb,
    scene::SceneInstance, transform::TransformSystem::TransformPropagate,
};
use bevy_dolly::prelude::*;
use bevy_inspector_egui::quick::WorldInspectorPlugin;
use bevy_polyline::prelude::*;
use leafwing_input_manager::prelude::*;
use structopt::StructOpt;

#[derive(Debug, StructOpt, Resource)]
#[structopt(name = "scene_aabb", about = "Draw AABB for a Bevy Scene")]
struct Opt {
    scene: String,
}

fn main() {
    App::new()
        .register_type::<SceneAabb>()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "scene_aabb".to_string(),
                ..default()
            }),
            ..default()
        }))
        .add_plugin(PolylinePlugin)
        .add_plugin(WorldInspectorPlugin::new())
        .add_plugin(InputManagerPlugin::<InputActions>::default())
        .add_system(Dolly::<MainCamera>::update_active)
        .add_startup_system(setup)
        .add_systems((handle_input, draw_scene_aabb, rotation_system))
        .add_system(
            create_scene_aabb
                .in_base_set(CoreSet::PostUpdate)
                .after(TransformPropagate),
        )
        .run();
}

fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
    let opt = Opt::from_args();

    // load the model
    let scene = asset_server.load(format!("{}#Scene0", opt.scene));
    commands.spawn((
        TargetScene,
        SceneBundle { scene, ..default() },
        InputManagerBundle::<InputActions> {
            action_state: ActionState::default(),
            input_map: input_map(),
        },
    ));

    // Add a camera
    commands.spawn((
        MainCamera,
        Camera3dBundle {
            tonemapping: Tonemapping::None,
            projection: OrthographicProjection {
                scaling_mode: bevy::render::camera::ScalingMode::WindowSize(48.0),
                ..default()
            }
            .into(),
            ..default()
        },
        InputManagerBundle::<InputActions> {
            action_state: ActionState::default(),
            input_map: input_map(),
        },
        Rig::builder()
            .with(Position::new(Vec3::new(0.0, 0.0, 0.0)))
            .with(YawPitch::new().pitch_degrees(-30.0).yaw_degrees(45.0))
            .with(Smooth::new_position(0.3))
            .with(Smooth::new_rotation(0.3))
            .with(Arm::new(Vec3::Z * 5.0))
            .build(),
    ));

    // Add a directional light (the sun)
    commands.spawn((DirectionalLightBundle {
        directional_light: DirectionalLight {
            illuminance: 18000.0,
            ..default()
        },
        transform: Transform::from_rotation(Quat::from_rotation_x(-0.5)),
        ..default()
    },));
}

#[derive(Component)]
struct MainCamera;

#[derive(Component)]
struct DebugBox;

#[derive(Component)]
pub struct TargetScene;

#[derive(Component)]
struct SceneAabbDebugBox(Entity);

#[derive(Component, Debug, Reflect)]
struct SceneAabb {
    min: Vec3,
    max: Vec3,
}

impl SceneAabb {
    fn new(center: Vec3) -> Self {
        Self {
            min: center,
            max: center,
        }
    }

    /// merge an child AABB into the Scene AABB
    fn merge_aabb(&mut self, aabb: &Aabb, global_transform: &GlobalTransform) {
        /*
              (2)-----(3)               Y
               | \     | \              |
               |  (1)-----(0) MAX       o---X
               |   |   |   |             \
          MIN (6)--|--(7)  |              Z
                 \ |     \ |
                  (5)-----(4)
        */

        let min = aabb.min();
        let max = aabb.max();
        let corners = [
            global_transform.transform_point(Vec3::new(max.x, max.y, max.z)),
            global_transform.transform_point(Vec3::new(min.x, max.y, max.z)),
            global_transform.transform_point(Vec3::new(min.x, max.y, min.z)),
            global_transform.transform_point(Vec3::new(max.x, max.y, min.z)),
            global_transform.transform_point(Vec3::new(max.x, min.y, max.z)),
            global_transform.transform_point(Vec3::new(min.x, min.y, max.z)),
            global_transform.transform_point(Vec3::new(min.x, min.y, min.z)),
            global_transform.transform_point(Vec3::new(max.x, min.y, min.z)),
        ];

        for corner in corners {
            let gt = corner.cmpgt(self.max);
            let lt = corner.cmplt(self.min);

            debug!("corner {:?}, gt {:?}, lt {:?}", corner, lt, gt);

            if gt.x {
                self.max.x = corner.x;
            } else if lt.x {
                self.min.x = corner.x;
            }

            if gt.y {
                self.max.y = corner.y;
            } else if lt.y {
                self.min.y = corner.y;
            }

            if gt.z {
                self.max.z = corner.z;
            } else if lt.z {
                self.min.z = corner.z;
            }
        }

        debug!("min {:?}, max {:?}", min, max);
    }
}

impl From<&SceneAabb> for Polyline {
    #[rustfmt::skip]
    fn from(scene_aabb: &SceneAabb) -> Polyline {
        /*
              (2)-----(3)               Y
               | \     | \              |
               |  (1)-----(0) MAX       o---X
               |   |   |   |             \
          MIN (6)--|--(7)  |              Z
                 \ |     \ |
                  (5)-----(4)
        */

        let min = scene_aabb.min;
        let max = scene_aabb.max;
        let corners = [
            Vec3::new(max.x, max.y, max.z), // 0
            Vec3::new(min.x, max.y, max.z), // 1
            Vec3::new(min.x, max.y, min.z), // 2
            Vec3::new(max.x, max.y, min.z), // 3
            Vec3::new(max.x, min.y, max.z), // 4
            Vec3::new(min.x, min.y, max.z), // 5
            Vec3::new(min.x, min.y, min.z), // 6
            Vec3::new(max.x, min.y, min.z), // 7
        ];

        Polyline {
            vertices: vec![
                corners[0], corners[1], corners[2], corners[3],
                corners[7], corners[4], corners[5], corners[6],
                corners[7], corners[3], corners[0], corners[4],
                corners[5], corners[1], corners[2], corners[6],
            ],
        }
    }
}

#[derive(Actionlike, PartialEq, Eq, Clone, Copy, Hash, Debug)]
pub enum InputActions {
    Click,
    Rotate,
    Scale,
    ResetCamera,
    ZeroCamera,
    RedoAabb,
    Run,
}

#[rustfmt::skip]
fn input_map() -> InputMap<InputActions> {
    InputMap::default()
        .insert(MouseButton::Left, InputActions::Click)
        .insert(DualAxis::mouse_motion(), InputActions::Rotate)
        .insert(SingleAxis::mouse_wheel_y(), InputActions::Scale)
        .insert(KeyCode::Z, InputActions::ResetCamera)
        .insert(KeyCode::Key0, InputActions::ZeroCamera)
        .insert(KeyCode::R, InputActions::RedoAabb)
        .insert(KeyCode::Space, InputActions::Run)
        .build()
}

fn handle_input(
    mut camera: Query<(&mut Rig, &mut Projection, &ActionState<InputActions>), With<MainCamera>>,
) {
    let (mut rig, mut projection, actions) = camera.single_mut();
    let camera_yp = rig.driver_mut::<YawPitch>();
    let Projection::Orthographic(projection) = projection.as_mut() else { panic!("wrong scaling mode") };

    if actions.just_pressed(InputActions::ResetCamera) {
        camera_yp.yaw_degrees = 45.0;
        camera_yp.pitch_degrees = -30.0;
        projection.scale = 1.0;
    }

    if actions.just_pressed(InputActions::ZeroCamera) {
        camera_yp.yaw_degrees = 0.0;
        camera_yp.pitch_degrees = 0.0;
        projection.scale = 1.0;
    }

    if actions.pressed(InputActions::Click) {
        let vector = actions.axis_pair(InputActions::Rotate).unwrap().xy();
        camera_yp.rotate_yaw_pitch(-0.1 * vector.x * 15.0, -0.1 * vector.y * 15.0);
    }

    let scale = actions.value(InputActions::Scale);
    if scale != 0.0 {
        projection.scale = (projection.scale * (1.0 - scale * 0.005)).clamp(0.001, 15.0);
    }
}

fn create_scene_aabb(
    mut commands: Commands,
    scene_instances: Query<(Entity, &SceneInstance, &GlobalTransform), Changed<GlobalTransform>>,
    scene_manager: Res<SceneSpawner>,
    children: Query<&Children>,
    bounding_boxes: Query<(&Aabb, &GlobalTransform)>,
) {
    for (entity, instance, global_transform) in scene_instances.iter() {
        if !scene_manager.instance_is_ready(**instance) {
            continue;
        }

        let mut scene_aabb = SceneAabb::new(global_transform.translation());

        for child in children.iter_descendants(entity) {
            let Ok((bb, transform)) = bounding_boxes.get(child) else { continue };
            scene_aabb.merge_aabb(bb, transform);
        }

        debug!("Scene Entity {:?}, AABB {:?}", entity, scene_aabb);
        commands.entity(entity).insert(scene_aabb);
    }
}

fn draw_scene_aabb(
    mut commands: Commands,
    scene_instances: Query<(Entity, &SceneAabb), Changed<SceneAabb>>,
    debug_boxes: Query<(Entity, &SceneAabbDebugBox)>,
    mut polyline_materials: ResMut<Assets<PolylineMaterial>>,
    mut polylines: ResMut<Assets<Polyline>>,
) {
    for (entity, aabb) in &scene_instances {
        for (debug_box, debug) in &debug_boxes {
            if debug.0 == entity {
                commands.entity(debug_box).despawn();
            }
        }

        commands.spawn((
            SceneAabbDebugBox(entity),
            PolylineBundle {
                polyline: polylines.add(aabb.into()),
                material: polyline_materials.add(PolylineMaterial {
                    width: 1.5,
                    color: Color::rgb(1.0, 0.0, 1.0),
                    perspective: false,
                    depth_bias: -0.0002,
                }),
                ..default()
            },
        ));
    }
}

/// Rotate the meshes to demonstrate how the bounding volumes update
fn rotation_system(time: Res<Time>, mut query: Query<&mut Transform, With<TargetScene>>) {
    for mut transform in query.iter_mut() {
        let rot_x = Quat::from_rotation_x((time.elapsed_seconds() / 5.0).sin() / 50.0);
        let rot_y = Quat::from_rotation_y((time.elapsed_seconds() / 3.0).sin() / 50.0);
        let rot_z = Quat::from_rotation_z((time.elapsed_seconds() / 4.0).sin() / 50.0);
        transform.rotate(rot_x * rot_y * rot_z);
    }
}

Thanks a lot. I was hoping someone already did this to save me some time, and they did :)