diff --git a/examples/widgets/painter.py b/examples/widgets/painter.py new file mode 100644 index 000000000000..930d2a38758d --- /dev/null +++ b/examples/widgets/painter.py @@ -0,0 +1,36 @@ +""" +=============================== +Painter Tool Demo +=============================== + +Drag the mouse to paint selected areas of the plot. The callback prints +the (x, y) coordinates of the center of the painted region. + +""" +from __future__ import print_function + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.widgets import Painter + +# Create a figure and axes for the plot. +fig, ax = plt.subplots() + + +# Generate data and create a scatter plot. +data = np.random.rand(100, 2) +data[:, 1] *= 2 +pts = ax.scatter(data[:, 0], data[:, 1], s=80) + + +# Define the "on_select" callback. +def callback(x, y): + print("(%3.2f, %3.2f)" % (x, y)) + + +print("\n click and drag \n (x, y)") + +# Create the painter tool and show the plot. +p = Painter(ax, callback) +p.label = 1 # set the color +plt.show() diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py index 2555037828f1..81f77ef59f11 100644 --- a/lib/matplotlib/tests/test_widgets.py +++ b/lib/matplotlib/tests/test_widgets.py @@ -271,3 +271,36 @@ def clicked_function(): pass cid = check.on_clicked(clicked_function) check.disconnect(cid) + + +@cleanup +def test_paint_tool(): + ax = get_ax() + + def onselect(x, y): + pass + + tool = widgets.Painter(ax, onselect) + + tool.radius = 10 + assert tool.radius == 10 + tool.label = 2 + assert tool.label == 2 + + do_event(tool, 'press', xdata=100, ydata=100) + do_event(tool, 'onmove', xdata=110, ydata=110) + do_event(tool, 'release') + + # Ensure that the correct number of pixels are labeled. + assert tool.overlay[tool.overlay == 2].size == 878 + + tool.label = 5 + do_event(tool, 'press', xdata=20, ydata=20) + do_event(tool, 'onmove', xdata=40, ydata=40) + do_event(tool, 'release') + + assert tool.overlay[tool.overlay == 5].size == 882 + assert tool.overlay[tool.overlay == 2].size == 878 + + tool.overlay = tool.overlay * 0 + assert tool.overlay.sum() == 0 diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py index a12aa82e83e5..0b4f04989994 100644 --- a/lib/matplotlib/widgets.py +++ b/lib/matplotlib/widgets.py @@ -19,6 +19,7 @@ import numpy as np from matplotlib import rcParams +from .colors import ListedColormap, NoNorm from .mlab import dist from .patches import Circle, Rectangle, Ellipse from .lines import Line2D @@ -1468,6 +1469,8 @@ def __init__(self, ax, onselect, useblit=False, button=None, if isinstance(button, int): self.validButtons = [button] + elif button is None: + self.validButtons = [1, 2, 3] else: self.validButtons = button @@ -2595,3 +2598,197 @@ def onmove(self, event): self.canvas.blit(self.ax.bbox) else: self.canvas.draw_idle() + + +LABELS_CMAP = ListedColormap(['white', 'red', 'dodgerblue', 'gold', + 'greenyellow', 'blueviolet']) + + +class Painter(_SelectorWidget): + """ + Paint regions on an axes. + + For the widget to remain responsive you must keep a reference to it. + + This is not meant to be used with polar coordinate plots. + + Resizing the plot will clear the existing selection. + + Example usage:: + + import numpy as np + import matplotlib.pyplot as plt + from matplotlib.widgets import Painter + + data = np.random.rand(100, 2) + data[:, 1] *= 2 + + fig, ax = plt.subplots() + pts = ax.scatter(data[:, 0], data[:, 1], s=80) + + def test(x, y): + print(x, y) + + p = Painter(ax, test) + plt.show() + """ + + def __init__(self, ax, on_select=None, overlay_props=None, + cursor_props=None, radius=5, cmap=LABELS_CMAP, + useblit=True, button=None, state_modifier_keys=None): + """ + Parameters + ========== + + ax: :class:`~matplotlib.axes.Axes` + The parent axes for the widget. + on_select: function + A callback for when a region is painted. Called with the + (x, y) coordinates of the region. + overlay_props: dict + The properties to apply to the overlay. + cursor_props: dcit + The properties to apply to the cursor. + radius: int + The radius of the cursor in pixels. + cmap: :class:~matplotlib.colorls.ListedColormap` + The colormap to use for the cursors. + useblit: bool + Whether to use blitting. + button: list + The button numbers supported for the tool (defaults to [1, 2, 3]) + state_modifier_keys: dict + A mapping of key names to state modifiers. + """ + super(Painter, self).__init__( + ax, on_select, useblit=useblit, button=button, + state_modifier_keys=None + ) + self.cmap = cmap + self._previous = None + self._overlay = None + self._overlay_plot = None + self._cursor_shape = [0, 0, 0] + + props = dict(edgecolor='r', facecolor='0.7', alpha=1, + animated=self.useblit, visible=False, zorder=2) + props.update(cursor_props or {}) + self._cursor = Rectangle((0, 0), 0, 0, **props) + self.ax.add_patch(self._cursor) + + x0, x1 = self.ax.get_xlim() + y0, y1 = self.ax.get_ylim() + if y0 < y1: + origin = 'lower' + else: + origin = 'upper' + props = dict(cmap=self.cmap, alpha=0.5, origin=origin, + norm=NoNorm(), visible=False, zorder=1, + extent=(x0, x1, y0, y1), aspect=self.ax.get_aspect()) + props.update(overlay_props or {}) + + self._resize(None) + self._overlay_plot = self.ax.imshow(self._overlay, **props) + self.artists = [self._cursor, self._overlay_plot] + self.connect_event('resize_event', self._resize) + + # These must be called last + self.label = 1 + self.radius = radius + self._drawing = True + for artist in self.artists: + artist.set_visible(True) + + @property + def overlay(self): + """The paint overlay image""" + return self._overlay + + @overlay.setter + def overlay(self, image): + self._overlay = image + if image is None: + self.ax.images.remove(self._overlay_plot) + self.update() + return + self._overlay_plot.set_data(image) + self._shape = image.shape + x0, x1 = self.ax.get_xlim() + y0, y1 = self.ax.get_ylim() + self._overlay_plot.set_extent((x0, x1, y0, y1)) + # Update the radii and window. + self.radius = self._radius + self.update() + + @property + def label(self): + """The label index""" + return self._label + + @label.setter + def label(self, value): + if value >= self.cmap.N: + raise ValueError('Maximum label value = %s' % len(self.cmap - 1)) + self._label = value + self._cursor.set_edgecolor(self.cmap(value)) + + @property + def radius(self): + """The label radius in pixels""" + return self._radius + + @radius.setter + def radius(self, r): + self._radius = r + xfm = self.ax.transData.inverted() + x0, y0 = xfm.transform((0, 0)) + x1, y1 = xfm.transform((r, r)) + self._rx, self._ry = abs(x1 - x0), abs(y1 - y0) + + self._cursor.set_width(self._rx * 2) + self._cursor.set_height(self._ry * 2) + + def _press(self, event): + self._update_cursor(event.xdata, event.ydata) + self._update_overlay(event.x, event.y) + self.update() + + def _onmove(self, event): + self._update_cursor(event.xdata, event.ydata) + if event.button and event.button in self.validButtons: + self._update_overlay(event.x, event.y) + self.update() + self.onselect(event.xdata, event.ydata) + + def _resize(self, event): + extents = self.ax.get_window_extent().extents + self._offsetx = extents[0] + self._offsety = extents[1] + self._shape = (int(extents[3] - extents[1]), + int(extents[2] - extents[0])) + self._overlay = np.zeros(self._shape, dtype='uint8') + if self._overlay_plot: + self._overlay_plot.set_data(self._overlay) + + def _release(self, event): + pass + + def _update_overlay(self, x, y): + col = x - self._offsetx + row = y - self._offsety + + h, w = self._shape + r = self._radius + + xmin = int(max(0, col - r)) + xmax = int(min(w, col + r + 1)) + ymin = int(max(0, row - r)) + ymax = int(min(h, row + r + 1)) + + self._overlay[slice(ymin, ymax), slice(xmin, xmax)] = self.label + self._overlay_plot.set_data(self._overlay) + + def _update_cursor(self, x, y): + x = x - self._rx + y = y - self._ry + self._cursor.set_xy((x, y))