通过套索或矩形选择进行选择。应该相当容易扩展。仅使用 matplotlib 小部件。
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, LassoSelector, RectangleSelector
from matplotlib.path import Path
from matplotlib.patches import Rectangle
from functools import partial
class LassoSelect(object):
def __init__(self, ax, collection):
self.canvas = ax.figure.canvas
self.collection = collection
self.xys = collection.get_offsets()
self.lasso = LassoSelector(ax, onselect=self.onselect)
self.setActive(False)
def onselect(self, verts):
path = Path(verts)
ind = np.nonzero([path.contains_point(xy) for xy in self.xys])[0]
fc = self.collection.get_facecolors()
fc[ind, -1] = 1
self.collection.set_facecolors(fc)
self.canvas.draw_idle()
def setActive(self, activate):
if activate:
self.lasso.active = True
else:
self.lasso.active = False
class RectangleSelect(object):
def __init__(self, ax, collection):
super(RectangleSelect, self).__init__()
self.RS = RectangleSelector(ax, self.onselect,
drawtype='box', useblit=True,
button=[1, 3], # don't use middle button
minspanx=5, minspany=5,
spancoords='pixels')
self.canvas = ax.figure.canvas
self.collection = collection
self.xys = collection.get_offsets()
self.setActive(False)
def onselect(self, eclick, erelease):
'eclick and erelease are the press and release events'
x1, y1 = eclick.xdata, eclick.ydata
x2, y2 = erelease.xdata, erelease.ydata
x0 = min(x1, x2)
y0 = min(y1, y2)
width = abs(x2-x1)
height = abs(y2-y1)
rect = Rectangle((x0, y0), width, height)
ind = np.nonzero([rect.contains_point(xy) for xy in self.xys])[0]
fc = self.collection.get_facecolors()
fc[ind, -1] = 1
self.collection.set_facecolors(fc)
self.canvas.draw_idle()
def setActive(self, activate):
if activate:
self.RS.set_active(True)
else:
self.RS.set_active(False)
class DataSelector(object):
activeColor = (0, 1, 0)
inActiveColor = (1, 1, 1)
def __init__(self, x, y):
super(DataSelector, self).__init__()
self.x, self.y = x, y
self.fig, self.ax = plt.subplots()
self.pts = self.ax.scatter(x, y)
self.selected_pts = []
self.addButtons()
self.selectorWidgets = {'Rectangle': RectangleSelect(self.ax, self.pts),
'Lasso': LassoSelect(self.ax, self.pts)}
self.fc = self.pts.get_facecolors()
if len(self.fc) == 0:
raise ValueError('Collection must have a facecolor')
elif len(self.fc) == 1:
self.fc = np.tile(self.fc, len(x)).reshape(len(x), -1)
self.fc[:, -1] = 0.3
self.pts.set_facecolors(self.fc)
def resetSelection(self, event=None):
self.fc[:, -1] = 0.3
self.fig.canvas.draw()
def exportSelection(self, event=None):
ind = np.nonzero(self.fc[:, -1] != 0.3)[0]
print self.x[ind], self.y[ind]
def activateWidget(self, key, event):
for widgetKey, widget in self.selectorWidgets.iteritems():
if widgetKey == key:
widget.setActive(True)
else:
widget.setActive(False)
for ax, button in self.buttonAxes.iteritems():
if ax == event.inaxes:
button.color = self.activeColor
else:
button.color = self.inActiveColor
button.hovercolor = button.color
def addButtons(self):
fig = self.fig
fig.subplots_adjust(bottom=0.25)
#Button text and callback
buttons = [("Lasso", None), ("Rectangle", None), ('Reset', self.resetSelection), ('Export Selection', self.exportSelection)]
nButtons = len(buttons)
axBorders = np.linspace(0.25, 0.90, nButtons+1, endpoint=False)
spacing = axBorders[1] - axBorders[0]
self.buttonAxes = {}
for i, btn in zip(axBorders, buttons):
#[left, bottom, width, height]
buttonAx = plt.axes([i, 0.1, spacing, 0.04])
button = Button(buttonAx, btn[0], color=self.inActiveColor)
if btn[1] is None:
button.on_clicked(partial(self.activateWidget, btn[0]))
else:
button.on_clicked(btn[1])
self.buttonAxes[buttonAx] = button
x, y = np.random.random((2, 25))*10
dataSelector = DataSelector(x, y)
plt.show()