From ed045de6880e8a410626d0f4b8db5a4830043cac Mon Sep 17 00:00:00 2001 From: Siqi Yan Date: Wed, 22 Nov 2023 19:23:04 -0700 Subject: [PATCH] added cacheable_coverage_progress.py --- nodes/cacheable_coverage_progress.py | 193 +++++++++++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100755 nodes/cacheable_coverage_progress.py diff --git a/nodes/cacheable_coverage_progress.py b/nodes/cacheable_coverage_progress.py new file mode 100755 index 0000000..2db7349 --- /dev/null +++ b/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