Выберите точки на графике python imshow() с помощью лассо или lassomanager

Может ли кто-нибудь помочь мне выбрать точки в цветном горшке 2D-массива, который я рисую с помощью imshow() из matplotlib. Я нашел два примера использования lasso или LassoSelector из matplotlib: http://matplotlib.org/examples/widgets/lasso_selector_demo.html http://matplotlib.org/examples/event_handling/lasso_demo.html

К сожалению, я не могу заставить их работать на меня с imshow и массивом 2d (обычно 50x500). Примеры работают нормально, но при использовании imshow python жалуется на такие вещи, как невозможность установить цвет. Я все еще довольно новичок в python. Насколько я понимаю, я должен получить коллекцию из данных, которые я отправляю в imshow? Ниже приведен пример кода, который не работает - я безуспешно пытался использовать оба примера, приведенные выше. Было бы здорово, если бы кто-нибудь указал мне правильное направление, хотя я тоже не против работать с кодом :) Спасибо.

import numpy as np
import matplotlib.pyplot as plt
# use random sample data:
random_data = np.random.rand(5, 5)
fig, ax = plt.subplots()
# Need the plot data as collection to be used with the lasso?
pts = ax.add_collection(imshow(data, aspect='auto', origin='lower',picker=True),autolim=False)
data = [Datum(*xy) for xy in random_data]
lman = ps.LassoManager(ax, data)

Вот код из примера, указанного выше:

from matplotlib.widgets import Lasso
from matplotlib.colors import colorConverter
from matplotlib.collections import RegularPolyCollection
from matplotlib import path

import matplotlib.pyplot as plt
from numpy import nonzero
from numpy.random import rand


class Datum(object):
    colorin = colorConverter.to_rgba('red')
    colorout = colorConverter.to_rgba('blue')

    def __init__(self, x, y, include=False):
        self.x = x
        self.y = y
        if include:
            self.color = self.colorin
        else:
            self.color = self.colorout


class LassoManager(object):
    def __init__(self, ax, data):
        self.axes = ax
        self.canvas = ax.figure.canvas
        self.data = data

        self.Nxy = len(data)

#        facecolors = [d.color for d in data]
        self.xys = [(d.x, d.y) for d in data]
        self.ind = []
        fig = ax.figure
        self.collection = RegularPolyCollection(
            fig.dpi, 6, sizes=(100,),
            facecolors=facecolors,
            offsets=self.xys,
            transOffset=ax.transData)

        ax.add_collection(self.collection)

        self.cid = self.canvas.mpl_connect('button_press_event', self.onpress)

    def callback(self, verts):
        facecolors = self.collection.get_facecolors()
        p = path.Path(verts)
        ind = p.contains_points(self.xys)
        self.ind = nonzero([p.contains_point(xy) for xy in self.xys])[0]
        for i in range(len(self.xys)):
            if ind[i]:
                facecolors[i] = colorConverter.to_rgba('red')
#                print ind
            else:
                facecolors[i] = colorConverter.to_rgba('blue')

        self.canvas.draw_idle()
        self.canvas.widgetlock.release(self.lasso)
        del self.lasso

    def onpress(self, event):
        if self.canvas.widgetlock.locked():
            return
        if event.inaxes is None:
            return
        self.lasso = Lasso(event.inaxes,
                           (event.xdata, event.ydata),
                           self.callback)
        # acquire a lock on the widget drawing
        self.canvas.widgetlock(self.lasso)

if __name__ == '__main__':
    print 'test'

    data = [Datum(*xy) for xy in rand(5, 5)]

    ax = plt.axes(xlim=(0, 1), ylim=(0, 1), autoscale_on=False)
    lman = LassoManager(ax, data)

    plt.show()

    for i in lman.ind:
        print lman.xys[i]

    print len(lman.xys)

Редактировать: приведенный ниже пример работает для меня, чтобы выбрать области произвольной формы из 2D-графика (не график разброса или xy). Не красиво и не элегантно, но это может подождать:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import LassoSelector
from matplotlib.path import Path


data = np.random.rand(10, 5)
grid = np.indices(data.shape)

ax = plt.subplot(111)
ax.pcolormesh(data)
ind = []

def onselect(verts):
    global ind
#    print verts
    pp = Path(verts)
    grid = [(i,j) for j in xrange(int(data.shape[0])) for i in xrange(int(data.shape[0]))]
    ii = np.nonzero([pp.contains_point(xy) for xy in grid])
    ind = [grid[i] for i in ii[0]]

#    ind = pp.contains_points(grid)
#    print ind
lasso = LassoSelector(ax, onselect)

plt.draw()
plt.show()

print ind