Skip to content

Commit

Permalink
Loading and saving model instances in the editor
Browse files Browse the repository at this point in the history
Signed-off-by: Michael X. Grey <[email protected]>
  • Loading branch information
mxgrey committed Jul 24, 2023
1 parent ec54d83 commit b9a2897
Show file tree
Hide file tree
Showing 18 changed files with 667 additions and 411 deletions.
73 changes: 69 additions & 4 deletions rmf_site_editor/src/site/load.rs
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,28 @@ fn generate_site_entities(commands: &mut Commands, site_data: &rmf_site_format::
.id();
id_to_entity.insert(*location_id, location);
consider_id(*location_id);

}

for (robot_id, robot) in &site_data.models.mobile_robots {
let robot_entity = site.spawn(robot.clone()).id();
id_to_entity.insert(*robot_id, robot_entity);
consider_id(*robot_id);
}

for (robot_id, robot) in &site_data.models.workcells {
let robot_entity = site.spawn(robot.clone()).id();
id_to_entity.insert(*robot_id, robot_entity);
consider_id(*robot_id);
}

for (scenario_id, scenario) in &site_data.scenarios {
let scenario_entity = site
.spawn(scenario.properties.clone())
.insert(SiteID(*scenario_id))
.id();
id_to_entity.insert(*scenario_id, scenario_entity);
consider_id(*scenario_id);
}
});

Expand All @@ -236,18 +258,61 @@ fn generate_site_entities(commands: &mut Commands, site_data: &rmf_site_format::
Ok(r) => r,
Err(id) => {
error!(
"ERROR: Nav Graph ranking could not load because a graph with \
"Nav Graph ranking could not load because a graph with \
id {id} does not exist."
);
RecencyRanking::new()
}
};

site_cmd
.insert(nav_graph_rankings)
.insert(NextSiteID(highest_id + 1));
site_cmd.insert(nav_graph_rankings);
let site_id = site_cmd.id();

// Construct instances separately so we can parent them correctly
for (scenario_id, scenario) in &site_data.scenarios {
let Some(scenario_entity) = id_to_entity.get(scenario_id).cloned() else {
error!(
"Failed to load scenario {scenario_id:?}, unable to find a \
loaded entity for it."
);
continue;
};

for (instance_id, instance) in &scenario.instances {
let Some(model_entity) = id_to_entity.get(&instance.model).cloned() else {
error!(
"Failed to load instance {instance_id:?} for scenario \
{scenario_id:?} because its model {:?} is missing.",
instance.model,
);
continue;
};

let Some(parent_entity) = id_to_entity.get(&instance.parent).cloned() else {
error!(
"Failed to load instance {instance_id:?} for scenario \
{scenario_id:?} because its parent {:?} is missing.",
instance.parent,
);
continue;
};

let instance_entity = commands
.spawn(instance.bundle.clone())
.insert(InScenario(scenario_entity))
.set_model_source(model_entity)
.insert(ModelMarker)
.insert(SiteID(*instance_id))
.id();
id_to_entity.insert(*instance_id, instance_entity);
consider_id(*instance_id);

commands.entity(parent_entity).add_child(instance_entity);
}
}

commands.entity(site_id).insert(NextSiteID(highest_id + 1));

// Make the lift cabin anchors that are used by doors subordinate
for (lift_id, lift_data) in &site_data.lifts {
for (_, door) in &lift_data.cabin_doors {
Expand Down
89 changes: 0 additions & 89 deletions rmf_site_editor/src/site/location.rs
Original file line number Diff line number Diff line change
Expand Up @@ -89,65 +89,6 @@ pub fn add_location_visuals(
}
}

pub fn add_robot_to_spawn_location(
mut commands: Commands,
locations: Query<
(
Entity,
&Point<Entity>,
&LocationTags,
Option<&LocationRobotModel>,
),
Changed<LocationTags>,
>,
levels: Query<(), With<LevelProperties>>,
mut models: Query<(&mut NameInSite, &mut AssetSource), With<ModelMarker>>,
anchors: AnchorParams,
parents: Query<&Parent>,
) {
for (e, point, location_tags, location_model) in &locations {
let position = anchors
.point_in_parent_frame_of(point.0, Category::Location, e)
.unwrap();

let parent = AncestorIter::new(&parents, point.0).find(|p| levels.get(*p).is_ok());

if let Some(parent) = parent {
if let Some(m) = location_tags.iter().find_map(|l| match l {
LocationTag::SpawnRobot(m) => Some(m),
_ => None,
}) {
if let Some(location_model) = location_model {
// Update existing model
if let Ok((mut name, mut source)) = models.get_mut(location_model.0) {
*name = m.name.clone();
*source = m.source.clone();
}
} else {
// Spawn new model
let mut model = m.clone();
model.pose = Pose {
trans: position.into(),
..default()
};
// TODO(luca) there should be a marker component to denote this is a robot
// as well as set it non static
// Robots should probably be made non deletable, since their spawning is
// controlled by the location property
let child = commands.spawn(model).id();
commands.entity(parent).push_children(&[child]);
commands.entity(e).insert(LocationRobotModel(child));
}
} else {
if let Some(location_model) = location_model {
commands.entity(location_model.0).despawn_recursive();
commands.entity(e).remove::<LocationRobotModel>();
}
}
}
}
}

pub fn update_changed_location(
mut locations: Query<
(
Expand Down Expand Up @@ -284,33 +225,3 @@ pub fn update_visibility_for_locations(
}
}
}

#[derive(Debug, Clone)]
pub struct ConsiderLocationTag {
pub tag: Option<LocationTag>,
pub for_element: Entity,
}

impl ConsiderLocationTag {
pub fn new(tag: Option<LocationTag>, for_element: Entity) -> Self {
Self { tag, for_element }
}
}

// TODO(MXG): Consider refactoring into a generic plugin, alongside ConsiderAssociatedGraph
pub fn handle_consider_location_tag(
mut recalls: Query<&mut RecallLocationTags>,
mut considerations: EventReader<ConsiderLocationTag>,
) {
for consider in considerations.iter() {
if let Ok(mut recall) = recalls.get_mut(consider.for_element) {
recall.consider_tag = consider.tag.clone();
let r = recall.as_mut();
if let Some(LocationTag::SpawnRobot(model)) | Some(LocationTag::Workcell(model)) =
&r.consider_tag
{
r.consider_tag_asset_source_recall.remember(&model.source);
}
}
}
}
7 changes: 4 additions & 3 deletions rmf_site_editor/src/site/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ pub use measurement::*;
pub mod model;
pub use model::*;

pub mod models;
pub use models::*;

pub mod nav_graph;
pub use nav_graph::*;

Expand Down Expand Up @@ -149,7 +152,7 @@ impl Plugin for SitePlugin {
.add_event::<ToggleLiftDoorAvailability>()
.add_event::<ExportLights>()
.add_event::<ConsiderAssociatedGraph>()
.add_event::<ConsiderLocationTag>()
.add_plugin(ModelSourcePlugin)
.add_plugin(ChangePlugin::<AssociatedGraphs<Entity>>::default())
.add_plugin(RecallPlugin::<RecallAssociatedGraphs<Entity>>::default())
.add_plugin(ChangePlugin::<Motion>::default())
Expand Down Expand Up @@ -234,7 +237,6 @@ impl Plugin for SitePlugin {
.with_system(update_floor_visibility)
.with_system(add_lane_visuals)
.with_system(add_location_visuals)
.with_system(add_robot_to_spawn_location)
.with_system(update_level_visibility)
.with_system(update_changed_lane)
.with_system(update_lane_for_moved_anchor)
Expand All @@ -249,7 +251,6 @@ impl Plugin for SitePlugin {
.with_system(update_changed_location)
.with_system(update_location_for_moved_anchors)
.with_system(handle_consider_associated_graph)
.with_system(handle_consider_location_tag)
.with_system(update_lift_for_moved_anchors)
.with_system(update_lift_door_availability)
.with_system(update_physical_lights)
Expand Down
152 changes: 152 additions & 0 deletions rmf_site_editor/src/site/models.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
/*
* Copyright (C) 2023 Open Source Robotics Foundation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/

use rmf_site_format::{AssetSource, Scale};
use bevy::{
prelude::*,
ecs::system::{Command, EntityCommands},
};

#[derive(Component, Clone, Copy)]
pub struct InScenario(pub Entity);

#[derive(Component, Clone, Copy)]
pub struct ModelSource(Entity);
impl ModelSource {
pub fn get(&self) -> Entity {
self.0
}
}

#[derive(Component, Clone)]
pub struct ModelInstances(Vec<Entity>);
impl ModelInstances {
pub fn iter(&self) -> impl Iterator<Item=&Entity> {
self.0.iter()
}
}

#[derive(Clone, Copy)]
struct ChangeModelSource {
instance: Entity,
source: Entity,
}

impl Command for ChangeModelSource {
fn write(self, world: &mut World) {
let mut previous_source = None;
if let Some(mut e) = world.get_entity_mut(self.instance) {
if let Some(previous) = e.get::<ModelSource>() {
if previous.0 == self.source {
// No need to change anything
return;
}

previous_source = Some(previous.0);
}
e.insert(ModelSource(self.source));
} else {
error!(
"Cannot change model source of instance {:?} because it no \
longer exists",
self.instance,
);
return;
}

if let Some(mut e) = world.get_entity_mut(self.source) {
if let Some(mut instances) = e.get_mut::<ModelInstances>() {
instances.0.push(self.instance);
} else {
e.insert(ModelInstances(vec![self.instance]));
}
} else {
error!(
"Cannot change model source of instance {:?} to {:?} because \
that source no longer exists",
self.instance,
self.source,
);

if let (Some(mut e), Some(prev)) = (
world.get_entity_mut(self.instance),
previous_source
) {
e.insert(ModelSource(prev));
}
return;
}

if let Some(previous) = previous_source {
if let Some(mut e) = world.get_entity_mut(previous) {
if let Some(mut instances) = e.get_mut::<ModelInstances>() {
instances.0.retain(|e| *e != self.instance);
}
}
}

world.send_event(self);
}
}

pub trait SetModelSourceExt {
fn set_model_source(&mut self, new_source: Entity) -> &mut Self;
}

impl<'w, 's, 'a> SetModelSourceExt for EntityCommands<'w, 's, 'a> {
fn set_model_source(&mut self, source: Entity) -> &mut Self {
let instance = self.id();
self.commands().add(ChangeModelSource { instance, source });
self
}
}

fn handle_changed_model_sources(
mut commands: Commands,
mut changed_instances: EventReader<ChangeModelSource>,
sources: Query<(&AssetSource, &Scale)>,
changed_sources: Query<
(&ModelInstances, &AssetSource, &Scale),
Or<(Changed<AssetSource>, Changed<Scale>)>,
>,
) {
for change in changed_instances.iter() {
let Some(mut instance) = commands.get_entity(change.source) else { continue };
let Ok((source, scale)) = sources.get(change.source) else { continue };
instance
.insert(source.clone())
.insert(scale.clone());
}

for (instances, source, scale) in &changed_sources {
for instance in instances.iter() {
let Some(mut instance) = commands.get_entity(*instance) else { continue };
instance
.insert(source.clone())
.insert(scale.clone());
}
}
}

pub struct ModelSourcePlugin;
impl Plugin for ModelSourcePlugin {
fn build(&self, app: &mut App) {
app
.add_event::<ChangeModelSource>()
.add_system(handle_changed_model_sources);
}
}
Loading

0 comments on commit b9a2897

Please sign in to comment.