Update SdfGenerator to save light data to file (#1209)

* save lights Signed-off-by: Ian Chen <[email protected]> * fix doc Signed-off-by: Ian Chen <[email protected]> * use updated api Signed-off-by: Ian Chen <[email protected]> * Update SdfGenerator to save joint data to file (#1220) * save joints Signed-off-by: Ian Chen <[email protected]> * use joinPaths Signed-off-by: Ian Chen <[email protected]> * Only output thread_pitch for screw joints Signed-off-by: Nate Koenig <[email protected]> Co-authored-by: Nate Koenig <[email protected]> Co-authored-by: Nate Koenig <[email protected]>
gazebosim · Nov 30, 2021 · 6dc70db · 6dc70db
1 parent 19f1a1a
commit 6dc70db
Show file tree

Hide file tree

Showing 4 changed files with 581 additions and 2 deletions.
diff --git a/src/SdfGenerator.cc b/src/SdfGenerator.cc
@@ -29,26 +29,33 @@
 #include "ignition/gazebo/components/AirPressureSensor.hh"
 #include "ignition/gazebo/components/Altimeter.hh"
 #include "ignition/gazebo/components/Camera.hh"
+#include "ignition/gazebo/components/ChildLinkName.hh"
 #include "ignition/gazebo/components/ContactSensor.hh"
 #include "ignition/gazebo/components/DepthCamera.hh"
 #include "ignition/gazebo/components/ForceTorque.hh"
 #include "ignition/gazebo/components/GpuLidar.hh"
 #include "ignition/gazebo/components/Imu.hh"
 #include "ignition/gazebo/components/Inertial.hh"
+#include "ignition/gazebo/components/Joint.hh"
+#include "ignition/gazebo/components/JointAxis.hh"
+#include "ignition/gazebo/components/JointType.hh"
 #include "ignition/gazebo/components/Light.hh"
 #include "ignition/gazebo/components/Link.hh"
 #include "ignition/gazebo/components/LogicalCamera.hh"
 #include "ignition/gazebo/components/Magnetometer.hh"
 #include "ignition/gazebo/components/Model.hh"
 #include "ignition/gazebo/components/Name.hh"
 #include "ignition/gazebo/components/ParentEntity.hh"
+#include "ignition/gazebo/components/ParentLinkName.hh"
 #include "ignition/gazebo/components/Pose.hh"
 #include "ignition/gazebo/components/RgbdCamera.hh"
 #include "ignition/gazebo/components/SelfCollide.hh"
+#include "ignition/gazebo/components/Sensor.hh"
 #include "ignition/gazebo/components/SourceFilePath.hh"
 #include "ignition/gazebo/components/SegmentationCamera.hh"
 #include "ignition/gazebo/components/Static.hh"
 #include "ignition/gazebo/components/ThermalCamera.hh"
+#include "ignition/gazebo/components/ThreadPitch.hh"
 #include "ignition/gazebo/components/WindMode.hh"
 #include "ignition/gazebo/components/World.hh"
 
@@ -295,7 +302,7 @@ namespace sdf_generator
     // First remove child entities of <world> whose names can be changed during
     // simulation (eg. models). Then we add them back from the data in the
     // ECM.
-    // TODO(addisu) Remove actors and lights
+    // TODO(addisu) Remove actors
     std::vector<sdf::ElementPtr> toRemove;
     if (_elem->HasElement("model"))
     {
@@ -305,13 +312,23 @@ namespace sdf_generator
         toRemove.push_back(modelElem);
       }
     }
+    if (_elem->HasElement("light"))
+    {
+      for (auto lightElem = _elem->GetElement("light"); lightElem;
+           lightElem = lightElem->GetNextElement("light"))
+      {
+        toRemove.push_back(lightElem);
+      }
+    }
+
     for (const auto &e : toRemove)
     {
       _elem->RemoveChild(e);
     }
 
     auto worldDir = common::parentPath(worldSdf->Data().Element()->FilePath());
 
+    // models
     _ecm.Each<components::Model, components::ModelSdf>(
         [&](const Entity &_modelEntity, const components::Model *,
             const components::ModelSdf *_modelSdf)
@@ -404,6 +421,21 @@ namespace sdf_generator
           return true;
         });
 
+    // lights
+    _ecm.Each<components::Light, components::ParentEntity>(
+        [&](const Entity &_lightEntity,
+            const components::Light *,
+            const components::ParentEntity *_parent) -> bool
+        {
+          if (_parent->Data() != _entity)
+            return true;
+
+           auto lightElem = _elem->AddElement("light");
+           updateLightElement(lightElem, _ecm, _lightEntity);
+
+          return true;
+        });
+
     return true;
   }
 
@@ -477,6 +509,21 @@ namespace sdf_generator
       }
     }
 
+    if (_elem->HasElement("joint"))
+    {
+      // update joints
+      sdf::ElementPtr jointElem = _elem->GetElement("joint");
+      while (jointElem)
+      {
+        std::string jointName = jointElem->Get<std::string>("name");
+        auto jointEnt = _ecm.EntityByComponents(
+            components::ParentEntity(_entity), components::Name(jointName));
+        if (jointEnt != kNullEntity)
+          updateJointElement(jointElem, _ecm, jointEnt);
+        jointElem = jointElem->GetNextElement("joint");
+      }
+    }
+
     return true;
   }
 
@@ -552,6 +599,21 @@ namespace sdf_generator
       }
     }
 
+    // update lights
+    if (_elem->HasElement("light"))
+    {
+      sdf::ElementPtr lightElem = _elem->GetElement("light");
+      while (lightElem)
+      {
+        std::string lightName = lightElem->Get<std::string>("name");
+        auto lightEnt = _ecm.EntityByComponents(
+            components::ParentEntity(_entity), components::Name(lightName));
+        if (lightEnt != kNullEntity)
+          updateLightElement(lightElem, _ecm, lightEnt);
+        lightElem = lightElem->GetNextElement("light");
+      }
+    }
+
     return true;
   }
 
@@ -563,7 +625,6 @@ namespace sdf_generator
     // Update sdf based on current components.
     // This list is to be updated as other components become updateable during
     // simulation
-
     auto updateSensorNameAndPose = [&]
     {
       // override name and pose sdf element using values from ECM
@@ -692,6 +753,183 @@ namespace sdf_generator
     return true;
   }
 
+  /////////////////////////////////////////////////
+  bool updateLightElement(sdf::ElementPtr _elem,
+                          const EntityComponentManager &_ecm,
+                          const Entity &_entity)
+  {
+    // Update sdf based on the light component
+    auto updateLightNameAndPose = [&]
+    {
+      // override name and pose sdf element using values from ECM
+      auto *nameComp = _ecm.Component<components::Name>(_entity);
+      _elem->GetAttribute("name")->Set(nameComp->Data());
+
+      auto *poseComp = _ecm.Component<components::Pose>(_entity);
+      auto poseElem = _elem->GetElement("pose");
+
+      // Remove all attributes of poseElem
+      for (const auto *attrName : {"relative_to", "degrees", "rotation_format"})
+      {
+        sdf::ParamPtr attr = poseElem->GetAttribute(attrName);
+        if (nullptr != attr)
+        {
+          attr->Reset();
+        }
+      }
+      poseElem->Set(poseComp->Data());
+      return true;
+    };
+
+    // light
+    auto lightComp = _ecm.Component<components::Light>(_entity);
+    if (lightComp)
+    {
+      const sdf::Light &light = lightComp->Data();
+      _elem->Copy(light.ToElement());
+      return updateLightNameAndPose();
+    }
+    return true;
+  }
+
+  /////////////////////////////////////////////////
+  bool updateJointElement(sdf::ElementPtr _elem,
+                          const EntityComponentManager &_ecm,
+                          const Entity &_entity)
+  {
+    // Update sdf based on the joint component
+    auto updateJointNameAndPose = [&]
+    {
+      // override name and pose sdf element using values from ECM
+      auto *nameComp = _ecm.Component<components::Name>(_entity);
+      _elem->GetAttribute("name")->Set(nameComp->Data());
+
+      auto *poseComp = _ecm.Component<components::Pose>(_entity);
+      auto poseElem = _elem->GetElement("pose");
+
+      // Remove all attributes of poseElem
+      for (const auto *attrName : {"relative_to", "degrees", "rotation_format"})
+      {
+        sdf::ParamPtr attr = poseElem->GetAttribute(attrName);
+        if (nullptr != attr)
+        {
+          attr->Reset();
+        }
+      }
+      poseElem->Set(poseComp->Data());
+      return true;
+    };
+
+    // joint
+    auto jointComp = _ecm.Component<components::Joint>(_entity);
+    if (!jointComp)
+    {
+      return false;
+    }
+
+    // joint type
+    auto jointTypeComp = _ecm.Component<components::JointType>(_entity);
+    sdf::JointType jointType = jointTypeComp->Data();
+    if (jointTypeComp)
+    {
+      std::string jointTypeStr = "invalid";
+      switch (jointType)
+      {
+        case sdf::JointType::BALL:
+          jointTypeStr = "ball";
+          break;
+        case sdf::JointType::CONTINUOUS:
+          jointTypeStr = "continuous";
+          break;
+        case sdf::JointType::FIXED:
+          jointTypeStr = "fixed";
+          break;
+        case sdf::JointType::PRISMATIC:
+          jointTypeStr = "prismatic";
+          break;
+        case sdf::JointType::GEARBOX:
+          jointTypeStr = "gearbox";
+          break;
+        case sdf::JointType::REVOLUTE:
+          jointTypeStr = "revolute";
+          break;
+        case sdf::JointType::REVOLUTE2:
+          jointTypeStr = "revolute2";
+          break;
+        case sdf::JointType::SCREW:
+          jointTypeStr = "screw";
+          break;
+        case sdf::JointType::UNIVERSAL:
+          jointTypeStr = "universal";
+          break;
+        default:
+          break;
+      }
+      _elem->GetAttribute("type")->Set<std::string>(jointTypeStr);
+    }
+
+    // parent
+    auto parentLinkNameComp =
+        _ecm.Component<components::ParentLinkName>(_entity);
+    if (parentLinkNameComp)
+    {
+      _elem->GetElement("parent")->Set<std::string>(parentLinkNameComp->Data());
+    }
+    // child
+    auto childLinkNameComp = _ecm.Component<components::ChildLinkName>(_entity);
+    if (childLinkNameComp)
+    {
+      _elem->GetElement("child")->Set<std::string>(childLinkNameComp->Data());
+    }
+    // thread pitch
+    auto threadPitchComp = _ecm.Component<components::ThreadPitch>(_entity);
+    if (threadPitchComp && jointType == sdf::JointType::SCREW)
+    {
+      _elem->GetElement("thread_pitch")->Set<double>(threadPitchComp->Data());
+    }
+    // axis
+    auto jointAxisComp = _ecm.Component<components::JointAxis>(_entity);
+    if (jointAxisComp)
+    {
+      const sdf::JointAxis axis = jointAxisComp->Data();
+      _elem->GetElement("axis")->Copy(axis.ToElement());
+    }
+    // axis2
+    auto jointAxis2Comp = _ecm.Component<components::JointAxis2>(_entity);
+    if (jointAxis2Comp)
+    {
+      const sdf::JointAxis axis2 = jointAxis2Comp->Data();
+      _elem->GetElement("axis2")->Copy(axis2.ToElement(1u));
+    }
+
+    // sensors
+    // remove existing ones in sdf element and add new ones from ECM.
+    std::vector<sdf::ElementPtr> toRemove;
+    if (_elem->HasElement("sensor"))
+    {
+      for (auto sensorElem = _elem->GetElement("sensor"); sensorElem;
+           sensorElem = sensorElem->GetNextElement("sensor"))
+      {
+        toRemove.push_back(sensorElem);
+      }
+    }
+    for (const auto &e : toRemove)
+    {
+      _elem->RemoveChild(e);
+    }
+
+    auto sensorEntities = _ecm.EntitiesByComponents(
+        components::ParentEntity(_entity), components::Sensor());
+
+    for (const auto &sensorEnt : sensorEntities)
+    {
+      sdf::ElementPtr sensorElem = _elem->AddElement("sensor");
+      updateSensorElement(sensorElem, _ecm, sensorEnt);
+    }
+
+    return updateJointNameAndPose();
+  }
+
   /////////////////////////////////////////////////
   /// \brief Checks if a string is a number
   /// \param[in] _str The string to check

diff --git a/src/SdfGenerator.hh b/src/SdfGenerator.hh
@@ -120,6 +120,28 @@ namespace sdf_generator
   bool updateSensorElement(sdf::ElementPtr _elem,
                            const EntityComponentManager &_ecm,
                            const Entity &_entity);
+
+  /// \brief Update an sdf::Element of a light.
+  /// Intended for internal use.
+  /// \input[in, out] _elem sdf::Element to update
+  /// \input[in] _ecm Immutable reference to the Entity Component Manager
+  /// \input[in] _entity Light entity
+  /// \returns true if update succeeded.
+  IGNITION_GAZEBO_VISIBLE
+  bool updateLightElement(sdf::ElementPtr _elem,
+                           const EntityComponentManager &_ecm,
+                           const Entity &_entity);
+
+  /// \brief Update an sdf::Element of a joint.
+  /// Intended for internal use.
+  /// \input[in, out] _elem sdf::Element to update
+  /// \input[in] _ecm Immutable reference to the Entity Component Manager
+  /// \input[in] _entity joint entity
+  /// \returns true if update succeeded.
+  IGNITION_GAZEBO_VISIBLE
+  bool updateJointElement(sdf::ElementPtr _elem,
+                           const EntityComponentManager &_ecm,
+                           const Entity &_entity);
 }  // namespace sdf_generator
 }  // namespace IGNITION_GAZEBO_VERSION_NAMESPACE
 }  // namespace gazebo