added cacheable_coverage_progress.py

nodes/cacheable_coverage_progress.py

@@ -0,0 +1,193 @@
+#! /usr/bin/env python3
+
+import rospy
+import tf2_ros
+from genpy import DeserializationError
+from nav_msgs.msg import OccupancyGrid
+from nav_msgs.srv import GetMap
+from numpy import ones, sum
+from std_msgs.msg import Float32, Header
+from std_srvs.srv import Trigger
+from tf.transformations import euler_from_quaternion
+
+# Constants for more readable index lookup
+X, Y, Z, W = 0, 1, 2, 3
+
+
+class CoverageProgressNode(object):
+    """
+    The CoverageProgressNode keeps track of coverage progress.
+    It does this by periodically looking up the position of the coverage disk in an occupancy grid.
+    Cells within a radius from this position are 'covered'
+
+    Cell values (0 ~ 100) are interpreted in this way: Lower is yet to be visited, Higher is covered
+    - 0: uncovered (initial value)
+    - > 0: covered
+
+    The node emits a coverage progress,
+    which is the ratio of cells considered coverage
+    (0.0 is everything uncovered, 1.0 is all covered)
+    """
+
+    DIRTY = 0
+    CLEAN = 100
+
+    def __init__(self):
+        # Transform Initialization
+        self.tfBuffer = tf2_ros.Buffer()
+        self.tfListener = tf2_ros.TransformListener(self.tfBuffer)
+
+        # ROS Stuff
+        self.map_frame = rospy.get_param("~map_frame", default="map")
+        self.coverage_frame = rospy.get_param("~coverage_frame", default="base_link")
+        self.cache_dir = rospy.get_param("~cache_dir", default="progress_map_data.txt")
+
+        self.progress_pub = rospy.Publisher("coverage_progress", Float32, queue_size=1)
+        self.grid_pub = rospy.Publisher("coverage_grid", OccupancyGrid, queue_size=1)
+
+        self.reset_srv = rospy.Service("reset", Trigger, self.reset)
+
+        # Receive Initial Static Map
+        self.grid = OccupancyGrid()
+        try:
+            with open(self.cache_dir, "rb") as f:
+                self.grid.deserialize(f.read())
+                self.grid.data = list(self.grid.data)
+        except DeserializationError:
+            rospy.logerr(
+                "Fail to load cached progress map. Used static map to initial a new progress map."
+            )
+            self._initialize_map()
+        except FileNotFoundError:
+            rospy.logwarn("Used static map to initial a new progress map.")
+            self._initialize_map()
+
+        # Get coverage radius, which is the robot radius which counts as coverage
+        try:
+            self.coverage_radius_meters = rospy.get_param("~coverage_radius")
+        except KeyError:
+            try:
+                self.coverage_radius_meters = rospy.get_param("~coverage_width") / 2.0
+            except KeyError:
+                rospy.logerr("Specify either coverage_width or coverage_radius")
+                raise ValueError(
+                    "Neither ~coverage_radius nor ~coverage_width specified, "
+                    "one of these is required"
+                )
+        self.coverage_radius_meters += (
+            2 * self.grid.info.resolution
+        )  # Compensate for discretization
+        self.coverage_radius_cells = int(
+            self.coverage_radius_meters / self.grid.info.resolution
+        )
+
+        # How covered is a cell after it has been covered for 1 time step
+        self.coverage_effectivity = rospy.get_param("~coverage_effectivity", default=5)
+
+        # Timer Callback
+        self._rate = rospy.get_param("~rate", 10.0)
+        self._update_timer = rospy.Timer(
+            rospy.Duration(1.0 / self._rate), self._update_callback
+        )
+        self._cache_timer = rospy.Timer(
+            rospy.Duration(secs=5), self._save_to_file
+        )
+
+    def _initialize_map(self):
+        print("Waiting for static map service...")
+        rospy.wait_for_service("static_map")
+        try:
+            static_map_client = rospy.ServiceProxy("static_map", GetMap)
+            response = static_map_client()
+            self.grid.header = response.map.header
+            self.grid.info = response.map.info
+            self.grid.data = list(response.map.data)
+
+            print("Received static map!")
+        except rospy.ServiceException as e:
+            print(f"static map service call failed: {e}")
+
+    def _update_callback(self, event):
+        # Get the position of point (0,0,0) the coverage_disk frame wrt.
+        # the map frame (which can both be remapped if need be)
+
+        try:
+            trans = self.tfBuffer.lookup_transform(
+                self.map_frame, self.coverage_frame, rospy.Time()
+            )
+        except (
+            tf2_ros.LookupException,
+            tf2_ros.ConnectivityException,
+            tf2_ros.ExtrapolationException,
+        ) as e:
+            print(f"tf error: {e}")
+            return
+
+        # Element of matrix corresponding to middle of coverage surface
+        x_point = int(
+            (trans.transform.translation.x - self.grid.info.origin.position.x)
+            / self.grid.info.resolution
+        )
+        y_point = int(
+            (trans.transform.translation.y - self.grid.info.origin.position.y)
+            / self.grid.info.resolution
+        )
+
+        # Initialize message
+        self.grid.header = Header()
+        self.grid.header.frame_id = self.map_frame
+
+        # Loop over amount of cells covered in x (j) and y (i) direction
+        for i in range(2 * self.coverage_radius_cells):
+            for j in range(2 * self.coverage_radius_cells):
+
+                # iterate from (-radius) to (+radius)
+                x_index = j - self.coverage_radius_cells
+                y_index = i - self.coverage_radius_cells
+
+                array_index = (
+                    x_point + x_index + self.grid.info.width * (y_point + y_index)
+                )
+
+                # if cell is inside coverage circle:
+                # determined using circle formula x^2+y^2 = (radius)^2
+                cell_in_coverage_circle = (
+                    x_index**2 + y_index**2 < self.coverage_radius_cells**2
+                )
+                # if cell is within the grid
+                cell_in_grid = 0 <= x_point + x_index < abs(
+                    int(self.grid.info.width / self.grid.info.resolution)
+                ) and 0 <= y_point + y_index < abs(
+                    int(self.grid.info.height / self.grid.info.resolution)
+                )
+
+                if cell_in_coverage_circle and cell_in_grid:
+                    self.grid.data[array_index] = min(
+                        self.CLEAN,
+                        self.grid.data[array_index] + self.coverage_effectivity,
+                    )
+
+        coverage_progress = float(sum(self.grid.data)) / (
+            self.grid.info.width * self.grid.info.height * 100.0
+        )
+
+        self.progress_pub.publish(coverage_progress)
+        self.grid_pub.publish(self.grid)
+
+    def reset(self, srv_request):
+        rospy.loginfo("Reset coverage progress and grid")
+        self._initialize_map()
+        return True, "Reset coverage progress and grid"
+
+    def _save_to_file(self, event):
+        with open(self.cache_dir, "wb") as f:
+            self.grid.serialize(f)
+
+
+if __name__ == "__main__":
+    rospy.init_node("coverage_progress")
+    try:
+        node = CoverageProgressNode()
+        rospy.spin()
+    except rospy.ROSInterruptException:
+        pass