From a0f053b27d0256be1f8203fb144da52dc0816f21 Mon Sep 17 00:00:00 2001
From: Sebastian Castro <sebas.a.castro@gmail.com>
Date: Thu, 23 Nov 2023 22:25:29 -0500
Subject: [PATCH] Remember original color of objects in planning scene

---
 .../moveit/planning_scene/planning_scene.h    | 16 +++++++++++
 .../planning_scene/src/planning_scene.cpp     | 27 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h b/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
index 472c63d68a..d8691fb6d8 100644
--- a/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
+++ b/moveit_core/planning_scene/include/moveit/planning_scene/planning_scene.h
@@ -53,6 +53,7 @@
 #include <octomap_msgs/msg/octomap_with_pose.hpp>
 #include <memory>
 #include <functional>
+#include <optional>
 #include <thread>
 #include <variant>
 #include <rclcpp/rclcpp.hpp>
@@ -741,7 +742,20 @@ class MOVEIT_PLANNING_SCENE_EXPORT PlanningScene : public std::enable_shared_fro
 
   bool hasObjectColor(const std::string& id) const;
 
+  /**
+   * \brief Gets the current color of an object.
+   * \param id The string id of the target object.
+   * \return The current object color.
+   */
   const std_msgs::msg::ColorRGBA& getObjectColor(const std::string& id) const;
+
+  /**
+   * \brief Tries to get the original color of an object, if one has been set before.
+   * \param id The string id of the target object.
+   * \return The original object color, if available, otherwise std::nullopt.
+   */
+  std::optional<std_msgs::msg::ColorRGBA> getOriginalObjectColor(const std::string& id) const;
+
   void setObjectColor(const std::string& id, const std_msgs::msg::ColorRGBA& color);
   void removeObjectColor(const std::string& id);
   void getKnownObjectColors(ObjectColorMap& kc) const;
@@ -1015,7 +1029,9 @@ class MOVEIT_PLANNING_SCENE_EXPORT PlanningScene : public std::enable_shared_fro
   StateFeasibilityFn state_feasibility_;
   MotionFeasibilityFn motion_feasibility_;
 
+  // Maps of current and original object colors (to manage attaching/detaching objects)
   std::unique_ptr<ObjectColorMap> object_colors_;
+  std::unique_ptr<ObjectColorMap> original_object_colors_;
 
   // a map of object types
   std::unique_ptr<ObjectTypeMap> object_types_;
diff --git a/moveit_core/planning_scene/src/planning_scene.cpp b/moveit_core/planning_scene/src/planning_scene.cpp
index 20d8364ff9..5a9aa502cd 100644
--- a/moveit_core/planning_scene/src/planning_scene.cpp
+++ b/moveit_core/planning_scene/src/planning_scene.cpp
@@ -1269,6 +1269,7 @@ bool PlanningScene::setPlanningSceneMsg(const moveit_msgs::msg::PlanningScene& s
   collision_detector_->cenv_->setPadding(scene_msg.link_padding);
   collision_detector_->cenv_->setScale(scene_msg.link_scale);
   object_colors_ = std::make_unique<ObjectColorMap>();
+  original_object_colors_ = std::make_unique<ObjectColorMap>();
   for (const moveit_msgs::msg::ObjectColor& object_color : scene_msg.object_colors)
     setObjectColor(object_color.id, object_color.color);
   world_->clearObjects();
@@ -1621,6 +1622,15 @@ bool PlanningScene::processAttachedCollisionObjectMsg(const moveit_msgs::msg::At
         const Eigen::Isometry3d& pose = attached_body->getGlobalPose();
         world_->addToObject(name, pose, attached_body->getShapes(), attached_body->getShapePoses());
         world_->setSubframesOfObject(name, attached_body->getSubframes());
+
+        // Try to set the object's color to its original color when first created.
+        // This ensures that the original color is reverted, e.g., when an object is attached and then unattached.
+        const auto original_object_color = getOriginalObjectColor(name);
+        if (original_object_color.has_value())
+        {
+          setObjectColor(attached_body->getName(), original_object_color.value());
+        }
+
         RCLCPP_DEBUG(LOGGER, "Detached object '%s' from link '%s' and added it back in the collision world",
                      name.c_str(), object.link_name.c_str());
       }
@@ -1975,6 +1985,17 @@ const std_msgs::msg::ColorRGBA& PlanningScene::getObjectColor(const std::string&
   return EMPTY;
 }
 
+std::optional<std_msgs::msg::ColorRGBA> PlanningScene::getOriginalObjectColor(const std::string& object_id) const
+{
+  if (original_object_colors_)
+  {
+    ObjectColorMap::const_iterator it = original_object_colors_->find(object_id);
+    if (it != original_object_colors_->end())
+      return it->second;
+  }
+  return std::nullopt;
+}
+
 void PlanningScene::getKnownObjectColors(ObjectColorMap& kc) const
 {
   kc.clear();
@@ -1997,6 +2018,12 @@ void PlanningScene::setObjectColor(const std::string& object_id, const std_msgs:
   if (!object_colors_)
     object_colors_ = std::make_unique<ObjectColorMap>();
   (*object_colors_)[object_id] = color;
+
+  // Set the original object color only once, if it's the first time adding this object ID.
+  if (!original_object_colors_)
+    original_object_colors_ = std::make_unique<ObjectColorMap>();
+  if (!getOriginalObjectColor(object_id))
+    (*original_object_colors_)[object_id] = color;
 }
 
 void PlanningScene::removeObjectColor(const std::string& object_id)