def __init__(self, map_name='fleming', bounds=[-5, -5, 5, 5], plot_robbers=True, map_display_type='probability', combined_only=True, publish_to_ROS=False): # Define map properties # <>TODO: Clean this up- add seperate map creation function? self.map_name = map_name if self.map_name == 'fleming': #self.bounds = [-9.5, -3.33, 4, 3.68] self.bounds = [0, 0, 5, 5] else: self.bounds = bounds # [x_min,y_min,x_max,y_max] in [m] self.plot_robbers = plot_robbers self.outer_bounds = [i * 1.1 for i in self.bounds] self.origin = [0, 0] # in [m] # <>TODO: Make display type relative to each robber self.display_type = map_display_type self.combined_only = combined_only # Set up ROS elements if using ROS self.publish_to_ROS = publish_to_ROS if publish_to_ROS: from cv_bridge import CvBridge import rospy from sensor_msgs.msg import Image from std_msgs.msg import String self.probability_target = 'Roy' self.bridge = CvBridge() self.image_pub = rospy.Publisher("python_probability_map", Image, queue_size=10) rospy.Subscriber("robot_probability_map", String, self.change_published_ax) # Define map elements self.objects = {} # For dynamic/static map objects (not robbers/cops) self.areas = {} self.cops = {} self.robbers = {} self.dynamic_elements = [] self.static_elements = [] self.information_elements = [] self.element_dict = {'dynamic': self.dynamic_elements, 'static': self.static_elements, 'information': self.information_elements} # Define layers self.shape_layer = ShapeLayer(self.element_dict, bounds=self.bounds) self.feasible_layer = FeasibleLayer(bounds=self.bounds) self.particle_layers = {} # One per robber, plus one combined self.probability_layers = {} # One per robber, plus one combined # Set up map if self.map_name == 'fleming': set_up_fleming(self) # <>TODO: make a generic 'setup map' function else: pass
class Map(object): """Environment map composed of multiple elements and layers. .. image:: img/classes_Map.png Parameters ---------- mapname : str The name of the map. bounds : array_like, optional Map boundaries as [xmin,ymin,xmax,ymax] in [m]. plot_robbers : list or bool A list of robbers to plot, True for all, False for None. map_display_type: {'particle', 'probability'} Defines which layer to use. combined_only : bool, optional Whether to show only the combined plot (as opposed to individual plots for each robber, plus one combined plot). Defaults to `True`. publish_to_ROS: bool Whether to publish an image topic of the map to ROS. """ # TODO: @Config Change plot_robbers to just be a string # TODO: @Refactor Seperate map and interface def __init__(self, map_name='fleming', bounds=[-5, -5, 5, 5], plot_robbers=True, map_display_type='probability', combined_only=True, publish_to_ROS=False): # Define map properties # <>TODO: Clean this up- add seperate map creation function? self.map_name = map_name if self.map_name == 'fleming': #self.bounds = [-9.5, -3.33, 4, 3.68] self.bounds = [0, 0, 5, 5] else: self.bounds = bounds # [x_min,y_min,x_max,y_max] in [m] self.plot_robbers = plot_robbers self.outer_bounds = [i * 1.1 for i in self.bounds] self.origin = [0, 0] # in [m] # <>TODO: Make display type relative to each robber self.display_type = map_display_type self.combined_only = combined_only # Set up ROS elements if using ROS self.publish_to_ROS = publish_to_ROS if publish_to_ROS: from cv_bridge import CvBridge import rospy from sensor_msgs.msg import Image from std_msgs.msg import String self.probability_target = 'Roy' self.bridge = CvBridge() self.image_pub = rospy.Publisher("python_probability_map", Image, queue_size=10) rospy.Subscriber("robot_probability_map", String, self.change_published_ax) # Define map elements self.objects = {} # For dynamic/static map objects (not robbers/cops) self.areas = {} self.cops = {} self.robbers = {} self.dynamic_elements = [] self.static_elements = [] self.information_elements = [] self.element_dict = {'dynamic': self.dynamic_elements, 'static': self.static_elements, 'information': self.information_elements} # Define layers self.shape_layer = ShapeLayer(self.element_dict, bounds=self.bounds) self.feasible_layer = FeasibleLayer(bounds=self.bounds) self.particle_layers = {} # One per robber, plus one combined self.probability_layers = {} # One per robber, plus one combined # Set up map if self.map_name == 'fleming': set_up_fleming(self) # <>TODO: make a generic 'setup map' function else: pass def add_obj(self, map_obj): self.objects[map_obj.name] = map_obj self.static_elements.append(map_obj) self.feasible_layer.define_feasible_regions(self.static_elements) def rem_obj(self, map_obj): self.static_elements.remove(map_obj) self.feasible_layer.define_feasible_regions(self.static_elements) del self.objects[map_obj.name] def add_robot(self, map_obj): # <>TODO: Modify so it can have relations self.dynamic_elements.append(map_obj) def rem_robot(self, map_obj): self.dynamic_elements.remove(map_obj) def add_area(self, area): self.areas[area.name] = area self.static_elements.append(area) def rem_area(self, area): self.static_elements.remove(area) del self.areas[area.name] def add_cop(self, cop_obj): self.dynamic_elements.append(cop_obj) self.cops[cop_obj.name] = cop_obj def rem_cop(self, cop_obj): self.dynamic_elements.remove(cop_obj) del self.cops[cop_obj.name] def add_robber(self, robber): # <>TODO: Make generic imaginary robbers if self.plot_robbers is True: robber.visible = True elif self.plot_robbers is False: robber.visible = False elif robber.name not in self.plot_robbers: robber.visible = False self.dynamic_elements.append(robber) self.robbers[robber.name] = robber def rem_robber(self, robber): robber.patch.remove() self.dynamic_elements.remove(robber) try: if self.fusion_engine is not None: if self.display_type == 'particle': self.rem_particle_layer(robber.name) elif self.display_type == 'probability': self.rem_probability_layer(robber.name) except: # <>TODO: actually catch other exceptions here logging.debug('No layer to remove.') del self.robbers[robber.name] def found_robber(self, robber): robber.visible = True robber.color = 'darkorange' try: if self.display_type == 'particle': self.rem_particle_layer(robber.name) elif self.display_type == 'probability': self.rem_probability_layer(robber.name) except: # <>TODO: actually catch other exceptions here logging.debug('No layer to remove.') def add_particle_layer(self, name, filter_): self.particle_layers[name] = ParticleLayer(filter_) def rem_particle_layer(self, name): self.particle_layers[name].remove() del self.particle_layers[name] def add_probability_layer(self, name, filter_): self.probability_layers[name] = ProbabilityLayer(filter_, fig=self.fig, bounds=self.bounds) def rem_probability_layer(self, name): self.probability_layers[name].remove() del self.probability_layers[name] def plot(self): # <>TODO: Needs rework fig = plt.figure(1, figsize=(12, 10)) ax = fig.add_subplot(111) self.shape_layer.update_plot(update_static=True) for area in self.areas.values(): ax.add_patch(area.get_patch()) ax.axis('scaled') ax.set_xlim([self.bounds[0], self.bounds[2]]) ax.set_ylim([self.bounds[1], self.bounds[3]]) ax.set_title('Experimental environment with landmarks and areas') ax.annotate('Kitchen', [-5, 2.5], weight='bold') ax.annotate('Billiard Room', [1.5, 2], weight='bold') ax.annotate('Library', [0.75, -2.25], weight='bold') ax.annotate('Study', [-4.5, -2.25], weight='bold') ax.annotate('Dining Room', [-8.75, -1.4], weight='bold') ax.annotate('Hallway', [-3.5, 0], weight='bold') plt.show() def setup_plot(self, fig=None, fusion_engine=None): """Create the initial plot for the animation. """ logging.info('Setting up plot') if fig is None: if plt.get_fignums(): self.fig = plt.gcf() else: self.fig = plt.figure(figsize=(14, 10)) else: self.fig = fig self.fusion_engine = fusion_engine self._setup_axes() self._setup_layers() def _setup_axes(self): self.axes = {} if len(self.robbers) == 1: name = self.robbers.iterkeys().next() self.axes[name] = self.fig.add_subplot(111) pos = self.axes[name].get_position() print pos pos = [pos.x0, pos.y0 * 1.2, pos.width, pos.height] print pos self.axes[name].set_position(pos) elif self.combined_only: self.axes['combined'] = self.fig.add_subplot(111) else: num_axes = len(self.robbers) + 1 num_rows = int(math.ceil(num_axes / 2)) i = 0 for robber_name in self.robbers: ax = plt.subplot2grid((num_rows, 4), (int(math.floor(i / 2)), (i % 2) * 2), colspan=2 ) self.axes[robber_name] = ax i += 1 # Add a plot for the combined estimate if (num_axes % 2) == 0: ax = plt.subplot2grid((num_rows, 4), (num_rows - 1, 2), colspan=2) else: ax = plt.subplot2grid((num_rows, 4), (num_rows - 1, 1), colspan=2) self.axes['combined'] = ax # Rescale, setup bounds and title for ax_name, ax in self.axes.iteritems(): ax.axis('scaled') ax.set_xlim([self.bounds[0], self.bounds[2]]) ax.set_xlabel('x position (m)') ax.set_ylim([self.bounds[1], self.bounds[3]]) ax.set_ylabel('y position (m)') if ax_name == 'combined': t = ax.set_title('Combined perception of all robots') else: t = ax.set_title("Map of {}'s perceived location" .format(ax_name)) try: if self.fusion_engine.vel_states is not None: t.set_y(1.2) except AttributeError: logging.debug('No vel states available.') # plt.tight_layout() def _setup_layers(self): # Set up basic layers self.shape_layers = {} self.feasible_layers = {} for ax_name, ax in self.axes.iteritems(): self.shape_layers[ax_name] = ShapeLayer(self.element_dict, bounds=self.bounds, ax=ax) # Set up probability/particle layers if self.fusion_engine is not None: filter_ = self.fusion_engine.filters[ax_name] self.probability_layers[ax_name] = \ ProbabilityLayer(filter_, fig=self.fig, ax=ax, bounds=self.bounds) def change_published_ax(self, msg): self.probability_target = msg.data def update(self, i=0): # self.shape_layer.update(i=i) for ax_name, ax in self.axes.iteritems(): try: self.shape_layers[ax_name].update(i=i) # Update probability/particle layers if self.fusion_engine is not None: if self.display_type == 'particle': self.particle_layers[ax_name].update(i=i) elif self.display_type == 'probability': self.probability_layers[ax_name].update(i=i) except KeyError: logging.debug('Robber already removed.') if self.publish_to_ROS and ax_name == self.probability_target and \ i % 1 == 0: import cv2 from cv_bridge import CvBridgeError extent = ax.get_window_extent().transformed( self.fig.dpi_scale_trans.inverted()) self.fig.savefig(ax_name + '.png', bbox_inches=extent.expanded(1.1, 1.2)) img = cv2.imread(ax_name + '.png', 1) try: self.image_pub.publish( self.bridge.cv2_to_imgmsg(img, "bgr8")) except CvBridgeError, e: print e