Ejemplo n.º 1
0
    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
Ejemplo n.º 2
0
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