diff --git a/rotorpy/utils/plotter.py b/rotorpy/utils/plotter.py index 0185497..d3e80e1 100644 --- a/rotorpy/utils/plotter.py +++ b/rotorpy/utils/plotter.py @@ -246,3 +246,87 @@ class Plotter(): self.estimator_exists = False return (time, x, x_des, v, v_des, q, q_des, w, s, s_des, M, T, wind, accel, gyro, accel_gt, x_mc, v_mc, q_mc, w_mc, filter_state, covariance, sd) + +def plot_map(ax, world_data, equal_aspect=True, color=None, edgecolor=None, alpha=1, world_bounds=True, axes=True): + """ + Plots the map in the world data in a top-down 2D view. + Inputs: + ax: The axis to plot on + world_data: The world data to plot + equal_aspect: Determines if the aspect ratio of the plot should be equal. + color: The color of the buildings. If None (default), it will use the color of the buildings. + edgecolor: The edge color of the buildings. If None (default), it will use the color of the buildings. + alpha: The alpha value of the buildings. If None (default), it will use the color of the buildings. + world_bounds: Whether or not to plot the world bounds as a dashed line around the 2D plot. + axes: Whether or not to plot the axis labels + Outputs: + Plots the map in the axis of interest. + """ + from matplotlib.patches import Rectangle + + # Add a dashed rectangle for the world bounds + if world_bounds: + world_patch = Rectangle((world_data['bounds']['extents'][0], world_data['bounds']['extents'][2]), + world_data['bounds']['extents'][1]-world_data['bounds']['extents'][0], world_data['bounds']['extents'][3]-world_data['bounds']['extents'][2], + linewidth=1, edgecolor='k', facecolor='none', linestyle='dashed') + ax.add_patch(world_patch) + + plot_xmin = world_data['bounds']['extents'][0] + plot_xmax = world_data['bounds']['extents'][1] + plot_ymin = world_data['bounds']['extents'][2] + plot_ymax = world_data['bounds']['extents'][3] + + for block in world_data['blocks']: + xmin = block['extents'][0] + xmax = block['extents'][1] + ymin = block['extents'][2] + ymax = block['extents'][3] + if color is None: + building_color = tuple(block['color']) + else: + building_color = color + if edgecolor is None: + building_edge_color = tuple(block['color']) + else: + building_edge_color = edgecolor + block_patch = Rectangle((xmin, ymin), (xmax-xmin), (ymax-ymin), linewidth=1, edgecolor=building_edge_color, facecolor=building_color, alpha=alpha, fill=True) + ax.add_patch(block_patch) + + if xmin < plot_xmin: + plot_xmin = xmin + if xmax > plot_xmax: + plot_xmax = xmax + if ymin < plot_ymin: + plot_ymin = ymin + if ymax > plot_ymax: + plot_ymax = ymax + + ax.set_xlim([plot_xmin, plot_xmax]) + ax.set_ylim([plot_ymin, plot_ymax]) + + if axes: + ax.set_xlabel("X (m)") + ax.set_ylabel("Y (m)") + + # Set the aspect ratio equal + if equal_aspect: + ax.set_aspect('equal') + + return + +if __name__ == "__main__": + + from rotorpy.world import World + + # Get a list of the maps available under worlds. + available_worlds = [fname for fname in os.listdir(os.path.abspath(os.path.join(os.path.dirname(__file__),'..','worlds'))) if 'json' in fname] + + # Load a random world + world_fname = np.random.choice(available_worlds) + world = World.from_file(os.path.abspath(os.path.join(os.path.dirname(__file__),'..','worlds', world_fname))) + + # Plot the world. + (fig, ax) = plt.subplots(nrows=1, ncols=1, num="Top Down World View") + plot_map(ax, world.world) + ax.set_title(world_fname) + plt.show() \ No newline at end of file