def draw_new_poly(self): coords = default_vertices(self.ax) Poly = Polygon(coords, animated=True, fc="white", ec='none', alpha=0.2, picker=True) self.polyList.append(Poly) self.ax.add_patch(Poly) x, y = zip(*Poly.xy) #line_color = 'none' color = np.array((1,1,1)) marker_face_color = (1,1,1) line_width = 4; line_kwargs = {'lw': line_width, 'color': color, 'mfc': marker_face_color} self.line.append(plt.Line2D(x, y, marker='o', alpha=1, animated=True, **line_kwargs)) self._update_line() self.ax.add_line(self.line[-1]) Poly.add_callback(self.poly_changed) Poly.num = self.poly_count self._ind = None # the active vert self.poly_count+=1;
class MaskDrawer(object): """An interactive polygon mask drawer on an image. Parameters ---------- ax : matplotlib plot axis Inpimg : 2d numpy array Input image to overlay for drawing mask Mask : Boolean numpy array same size of inpimg A Mask which will be used as initial mask and updated upon confirmation max_ds : float Max pixel distance to count as a vertex hit. PolyAtStart : List of vertices A list of square vertices to draw the initial polygon Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices 'n' : Invert the region selected by polynomial for masking 'c' : Confirm the polygon and update the mask """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit def __init__(self, ax, Inpimg, Mask, max_ds=10,PolyAtStart = [(50,50),(100,50),(100,100),(50,100)]): self.showverts = True self.max_ds = max_ds self.Mask = Mask self.img = Inpimg self.maskinvert = False # imshow the image self.imgplot = ax.imshow(np.ma.filled(np.ma.array(self.img,mask=self.Mask,fill_value=np.nan)), cmap=cm.gray) self.poly = Polygon(PolyAtStart, animated=True, fc='y', ec='none', alpha=0.5) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title("Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "'n' to invert the region for masking, 'c' to confirm & apply the mask.") self.ax = ax x, y = zip(*self.poly.xy) self.line = Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.7, animated=True) # self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind]>=self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes==None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='t': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key=='d': ind = self.get_ind_under_point(event) if ind is not None: self.poly.xy = [tup for i,tup in enumerate(self.poly.xy) if i!=ind] self.line.set_data(zip(*self.poly.xy)) elif event.key=='i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # display coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d<=self.epsilon: self.poly.xy = np.array( list(self.poly.xy[:i]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i:])) self.line.set_data(zip(*self.poly.xy)) break elif event.key=='n': """ Flips the region inside out of the polygon to be masked """ print('Inverting the mask region') self.maskinvert = not self.maskinvert elif event.key=='c': """ Confirm the drawn polynomial shape and add update the mask """ self.UpdateMask() #Update the imshowed image with new mask self.imgplot.set_data(np.ma.filled(np.ma.array(self.img,mask=self.Mask,fill_value=np.nan))) self.imgplot.figure.canvas.draw() self.canvas.draw() def UpdateMask(self): """ Updates the maks array with points insied the polygon """ print('Updating the original Mask..') Path = self.poly.get_path() h, w = self.Mask.shape y, x = np.mgrid[:h,:w] XYpoints = np.transpose((x.ravel(), y.ravel())) NewMask = Path.contains_points(XYpoints) if self.maskinvert : NewMask = ~NewMask # Combine the two mask by taing an elemet wise or self.Mask = self.Mask | NewMask.reshape((h,w)) def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x,y = event.xdata, event.ydata self.poly.xy[self._ind] = x,y self.line.set_data(zip(*self.poly.xy)) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox)
class LineInteractor(QObject): """ A Gaussian line interactor. Arguments: ax axes to which the interactor is associated c0 intercept of continuum cs slope of continuum x0 center of the line A amplitude of the line fwhm FWHM of the line epsilon max pixel distance to count as a vertex hit Key-bindings 't' toggle markers on and off. When markers are on, one can move or modify the line 'd' delete the line """ showverts = True mySignal = pyqtSignal(str) modSignal = pyqtSignal(str) def __init__(self, ax, c0, cs, x0, A, fwhm, color='#7ec0ee', epsilon=10): super().__init__() sys.setrecursionlimit(10000) # 10000 is 10x the default value self.epsilon = epsilon self.ax = ax canvas = ax.figure.canvas self.c0 = c0 # Value of continuum at origin self.cs = cs # Slope of the continuum self.type = 'Line' self.color = color self.x0 = x0 self.A = A self.fwhm = fwhm self.computeMarkers() self.computeGaussian() self.poly = Polygon(self.verts, animated=True, fill=False, closed=False, color=self.color) self.ax.add_patch(self.poly) x, y = zip(*self.xy) self.line = Line2D(x, y, marker='o', linestyle=None, linewidth=0., markerfacecolor=color, animated=True) self.ax.add_line(self.line) self.artists = [self.poly, self.line] self.cid = self.poly.add_callback(self.poly_changed) self._ind = None # the active vert #self.cid_draw = canvas.mpl_connect('draw_event', self.draw_callback) self.cid_press = canvas.mpl_connect('button_press_event', self.button_press_callback) self.cid_key = canvas.mpl_connect('key_press_event', self.key_press_callback) self.cid_release = canvas.mpl_connect('button_release_event', self.button_release_callback) #self.cid_motion = canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def computeMarkers(self): 'Compute position of markers.' x = self.x0 + 0.5 * self.fwhm * np.array([-1, 0, 1]) y = self.c0 + (x - self.x0) * self.cs + self.A * np.array( [0.5, 1., 0.5]) self.xy = [(i, j) for (i, j) in zip(x, y)] def computeGaussian(self): 'Compute the Gaussian polygon from the position of the markers.' self.sigma = self.fwhm / (2 * np.sqrt(2 * np.log(2))) # Create an array of x values and compute the value of the Gaussian on it x = np.linspace(self.x0 - self.fwhm, self.x0 + self.fwhm, 30) dx = (x - self.x0) / self.sigma / np.sqrt(2.) y = self.c0 + (x - self.x0) * self.cs + self.A * np.exp(-dx * dx) self.verts = [(x_, y_) for x_, y_ in zip(x, y)] def draw_callback(self, event): #self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # Distance is computed in pixels on the screen xy = self.ax.transData.transform(self.xy) x, y = zip(*xy) x = np.array(x) y = np.array(y) d = np.hypot(x - event.x, y - event.y) indseq, = np.nonzero(d == d.min()) ind = indseq[0] if d[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes is None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None # Notify callback self.modSignal.emit('line guess modified') def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_point(event) if ind is not None: self.mySignal.emit('line deleted ') self.canvas.draw_idle() def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x_, y_ = event.xdata, event.ydata x, y = zip(*self.xy) if x[-1] > x[0]: # case wavelength if self._ind == 0: if x_ < x[1]: self.fwhm = 2 * (x[1] - x_) elif self._ind == 1: dx = x_ - x[1] self.x0 += dx self.c0 += dx * self.cs dy = y_ - y[1] if (self.A > 0) & (dy < -self.A): # Emission line pass elif (self.A < 0) & (dy > -self.A): # Absorption line pass else: self.A += dy elif self._ind == 2: if x_ > x[1]: self.fwhm = 2 * (x_ - x[1]) else: if self._ind == 0: if x_ > x[1]: self.fwhm = 2 * (x_ - x[1]) elif self._ind == 1: dx = x_ - x[1] self.x0 += dx self.c0 += dx * self.cs dy = y_ - y[1] if (self.A > 0) & (dy < -self.A): # Emission line pass elif (self.A < 0) & (dy > -self.A): # Absorption line pass else: self.A += dy elif self._ind == 2: if x_ < x[1]: self.fwhm = 2 * (x[1] - x_) # Update and redraw self.updateCurves() def updateCurves(self): self.computeGaussian() self.computeMarkers() self.line.set_data(zip(*self.xy)) self.poly.xy = self.verts #self.canvas.restore_region(self.background) #self.ax.draw_artist(self.poly) #self.ax.draw_artist(self.line) #self.canvas.update() #self.canvas.flush_events() def disconnect(self): #self.canvas.mpl_disconnect(self.cid_draw) self.canvas.mpl_disconnect(self.cid_press) self.canvas.mpl_disconnect(self.cid_release) #self.canvas.mpl_disconnect(self.cid_motion) self.canvas.mpl_disconnect(self.cid_key) try: self.line.remove() except BaseException: print('no markers') try: self.gauss.remove() except BaseException: print('no line') self.canvas.draw_idle()
class PolygonInteractor(QObject): """ An polygon editor. Key-bindings 't' toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' delete the vertex under point 'i' insert a vertex at point. You must be within epsilon of the line connecting two existing vertices """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit mySignal = pyqtSignal(str) modSignal = pyqtSignal(str) def __init__(self, ax, verts): super().__init__() from matplotlib.patches import Polygon from matplotlib.lines import Line2D # from matplotlib.artist import Artist self.ax = ax self.type = 'Polygon' self.poly = Polygon(list(verts), animated=True, fill=False, closed=True, color='lime') self.ax.add_patch(self.poly) self.canvas = self.poly.figure.canvas x, y = zip(*self.poly.xy) self.line = Line2D(x, y, marker='o', linestyle=None, linewidth=0., markerfacecolor='g', animated=True) self.ax.add_line(self.line) self.cid = self.poly.add_callback(self.poly_changed) self._ind = None # the active vert self.connect() self.aperture = self.poly def connect(self): self.cid_draw = self.canvas.mpl_connect('draw_event', self.draw_callback) self.cid_press = self.canvas.mpl_connect('button_press_event', self.button_press_callback) self.cid_key = self.canvas.mpl_connect('key_press_event', self.key_press_callback) self.cid_release = self.canvas.mpl_connect('button_release_event', self.button_release_callback) self.cid_motion = self.canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas.draw_idle() def disconnect(self): self.canvas.mpl_disconnect(self.cid_draw) self.canvas.mpl_disconnect(self.cid_press) self.canvas.mpl_disconnect(self.cid_key) self.canvas.mpl_disconnect(self.cid_release) self.canvas.mpl_disconnect(self.cid_motion) self.poly.remove() self.line.remove() self.canvas.draw_idle() self.aperture = None def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) #self.canvas.blit(self.ax.bbox) #self.canvas.udpate() #self.canvas.flush_events() def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.hypot(xt - event.x, yt - event.y) indseq, = np.nonzero(d == d.min()) ind = indseq[0] if d[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes is None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_point(event) if ind is not None: if len(self.poly.xy) < 5: # the minimum polygon has 4 points since the 1st is repeated as final # Delete polygon #self.disconnect() #self.poly = None #self.line = None self.mySignal.emit('polygon deleted') else: self.poly.xy = [tup for i, tup in enumerate(self.poly.xy) if i != ind] self.line.set_data(zip(*self.poly.xy)) self.mySignal.emit('one vertex of polygon removed') elif event.key == 'i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # display coords for i in range(len(xys) - 1): s0 = xys[i] s1 = xys[i + 1] d = dist_point_to_segment(p, s0, s1) if d <= self.epsilon: self.poly.xy = np.array( list(self.poly.xy[:i+1]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i+1:])) self.line.set_data(zip(*self.poly.xy)) break self.canvas.draw_idle() def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x, y = event.xdata, event.ydata self.poly.xy[self._ind] = x, y if self._ind == 0: self.poly.xy[-1] = x, y elif self._ind == len(self.poly.xy) - 1: self.poly.xy[0] = x, y self.updateMarkers() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() self.canvas.flush_events() # Notify callback self.modSignal.emit('polygon modified') def updateMarkers(self): self.line.set_data(zip(*self.poly.xy))
class MaskCreator(object): """An interactive polygon editor. Parameters ---------- poly_xy : list of (float, float) List of (x, y) coordinates used as vertices of the polygon. max_ds : float Max pixel distance to count as a vertex hit. Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices """ def __init__(self, ax, poly_xy=None, max_ds=10): self.showverts = True self.max_ds = max_ds if poly_xy is None: poly_xy = default_vertices(ax) self.poly = Polygon(poly_xy, animated=True, fc='y', ec='none', alpha=0.4) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title("Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "Close figure when done.") self.ax = ax x, y = zip(*self.poly.xy) self.line = plt.Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.2, animated=True) self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def get_mask(self, shape): """Return image mask given by mask creator""" h, w = shape y, x = np.mgrid[:h, :w] points = np.transpose((x.ravel(), y.ravel())) #mask = nxutils.points_inside_poly(points, self.verts) #mask = Path(points).contains_points(self.verts) path = Path(self.verts) mask = path.contains_points(points) return mask.reshape(h, w) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() #Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def button_press_callback(self, event): 'whenever a mouse button is pressed' ignore = not self.showverts or event.inaxes is None or event.button != 1 if ignore: return self._ind = self.get_ind_under_cursor(event) def button_release_callback(self, event): 'whenever a mouse button is released' ignore = not self.showverts or event.button != 1 if ignore: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_cursor(event) if ind is None: return if ind == 0 or ind == self.last_vert_ind: print("Cannot delete root node") return self.poly.xy = [ tup for i, tup in enumerate(self.poly.xy) if i != ind ] self._update_line() elif event.key == 'i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # cursor coords for i in range(len(xys) - 1): s0 = xys[i] s1 = xys[i + 1] d = my_dist_point_to_segment(p, s0, s1) if d <= self.max_ds: self.poly.xy = np.array( list(self.poly.xy[:i + 1]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i + 1:])) self._update_line() break self.canvas.draw() def motion_notify_callback(self, event): 'on mouse movement' ignore = (not self.showverts or event.inaxes is None or event.button != 1 or self._ind is None) if ignore: return x, y = event.xdata, event.ydata if self._ind == 0 or self._ind == self.last_vert_ind: self.poly.xy[0] = x, y self.poly.xy[self.last_vert_ind] = x, y else: self.poly.xy[self._ind] = x, y self._update_line() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def _update_line(self): # save verts because polygon gets deleted when figure is closed self.verts = self.poly.xy self.last_vert_ind = len(self.poly.xy) - 1 self.line.set_data(zip(*self.poly.xy)) def get_ind_under_cursor(self, event): 'get the index of the vertex under cursor if within max_ds tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.max_ds: ind = None return ind
class PolyClick(object): """ p = PolyClick() # Start interactive polygon generator, or p = PolyClick(verts, ax=gca()) # Initialize polygon based on verts, or p = PolyClick(x, y, ax=gca()) # Initialize polygon based on x and y, or p = PolyClick('verts.pickle', ax=gca()) # Load in verts from pickle file If verts or (x, y) are not given, an interactive polygon creation session is started. Verticies are orange. Switch to editing mode by hitting return. This changes the vertecies to black. The points may be moved with the mouse, or modified with the key commands listed below. Current data are always available in bry.x, bry.y and bry.verts. Key commands: enter : switch to grid editing mode t : toggle visibility of verticies d : delete a vertex i : insert a vertex at a point on the polygon line Methods: p.dump(bry_file='bry.pickle) Write the current boundary informtion (bry.verts) to a cPickle file bry_file. p.load(bry_file='bry.pickle) Read in boundary informtion (verts) from the cPickle file bry_file. Attributes: p.x, p.y : The x and y locations of the polygon verticies. p.verts : The verticies of the polygon (equiv. to zip(x, y)) p.area : the area of the polygon. Positive for a counterclockwise path, negative for a clockwise path. p.centroid : the (x, y) location of the polygon centroid. """ _showverts = True _epsilon = 5 # max pixel distance to count as a vertex hit def _init_boundary_interactor(self): """Send polyclick thread to boundary interactor""" # Get rid old mpl connections. self._ax.figure.canvas.mpl_disconnect(self._key_id) self._ax.figure.canvas.mpl_disconnect(self._click_id) # Shade the selected region in a polygon self._poly = Polygon(self.verts, alpha=0.2, fc='k', animated=True) self._ax.add_patch(self._poly) # modify the line properties self._line.set_markerfacecolor('k') # get the canvas and connect the callback events cid = self._poly.add_callback(self._poly_changed) self._ind = None # the active vert self._canvas.mpl_connect('draw_event', self._draw_callback) self._canvas.mpl_connect('button_press_event', self._button_press_callback) self._canvas.mpl_connect('key_press_event', self._key_press_callback) self._canvas.mpl_connect('button_release_event', self._button_release_callback) self._canvas.mpl_connect('motion_notify_event', self._motion_notify_callback) def _draw_callback(self, event): self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self._line.get_visible() Artist.update_from(self._line, poly) self._line.set_visible(vis) # don't use the poly visibility state def _get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' try: x, y = zip(*self._poly.xy) # display coords xt, yt = self._poly.get_transform().numerix_x_y(x, y) d = sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = nonzero(equal(d, amin(d))) ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind except: # display coords xy = asarray(self._poly.xy) xyt = self._poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = nonzero(equal(d, amin(d)))[0] ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind def _button_press_callback(self, event): 'whenever a mouse button is pressed' # if not self._showverts: return if event.inaxes==None: return if event.button != 1: return self._ind = self._get_ind_under_point(event) def _button_release_callback(self, event): 'whenever a mouse button is released' # if not self._showverts: return if event.button != 1: return self._ind = None def _key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='t': self._showverts = not self._showverts elif event.key=='d': ind = self._get_ind_under_point(event) if ind is not None: self._poly.xy = [tup for i,tup in enumerate(self._poly.xy) if i!=ind] self._line.set_data(zip(*self._poly.xy)) elif event.key=='i': xys = self._poly.get_transform().seq_xy_tups(self._poly.xy) p = event.x, event.y # display coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy.insert(i+1, (event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) break s0 = xys[-1] s1 = xys[0] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy.append((event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) self._draw_callback(event) self._canvas.draw() def _motion_notify_callback(self, event): 'on mouse movement' if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x,y = event.xdata, event.ydata self._poly.xy[self._ind] = x,y x, y = zip(*self._poly.xy) self._line.set_data(x, y) self._canvas.restore_region(self._background) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _on_return(self, event): if event.key in ('enter', None): self._init_boundary_interactor() def _on_click(self, event): self.x.append(event.xdata) self.y.append(event.ydata) self._line.set_data(self.x, self.y) self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def __init__(self, *args, **kwargs): ax = kwargs.pop('ax', None) if ax is None: ax = pl.gca() self._ax = ax assert len(args) <=2, 'Input is either verticies (1 arg), or x and y (2 args)' if args is not (): if isinstance(args[0], str): verts = load(x) x, y = zip(*verts) x = list(x); y = list(y) elif len(args) == 1: # assume verticies input verts = args[0] x, y = zip(*verts) x = list(x); y = list(y) else: # x and y inputs x = list(args[0]) y = list(args[1]) assert len(x) == len(y), 'x and y must have the same length.' else: # no input x = []; y = [] self._line = pl.Line2D(x, y, marker='o', markerfacecolor='orange', animated=True) self._ax.add_line(self._line) self._canvas = self._line.figure.canvas self._key_id = self._ax.figure.canvas.mpl_connect('key_press_event', self._on_return) self._click_id = self._ax.figure.canvas.mpl_connect('button_press_event', self._on_click) if len(x) > 0: self._init_boundary_interactor() def dump(self, poly_file): f = open(bry_file, 'wb') cPickle.dump(self.verts, f, protocol=-1) f.close() def load(self, poly_file): verts = load(bry_file) x, y = zip(*verts) self._line.set_data(x, y) if hasattr(self, '_poly'): self._poly.xy = verts self._draw_callback(None) self._canvas.draw() def inside(self, *args): """docstring for inside""" assert len(args) > 0, 'must provide either verticies or x and y.' if len(args) == 1: # assume verticies input verts = args[0] else: # x and y inputs x = list(args[0]) y = list(args[1]) assert len(x) == len(y), 'x and y must have the same length.' verts = zip(x, y) return self._geom.inside(verts) def get_verts(self):return zip(self.x, self.y) verts = property(get_verts) def get_xdata(self): return self._line.get_xdata() x = property(get_xdata) def get_ydata(self): return self._line.get_ydata() y = property(get_ydata) def _get_geom(self): return Polygeom(self.verts) _geom = property(_get_geom) def _get_area(self): return self._geom.area area = property(_get_area) def _get_centroid(self): return self._geom.centroid centroid = property(_get_centroid)
class PolyClick(object): """ p = PolyClick() # Start interactive polygon generator, or p = PolyClick(verts, ax=gca()) # Initialize polygon based on verts, or p = PolyClick(x, y, ax=gca()) # Initialize polygon based on x and y, or p = PolyClick('verts.pickle', ax=gca()) # Load in verts from pickle file If verts or (x, y) are not given, an interactive polygon creation session is started. Verticies are orange. Switch to editing mode by hitting return. This changes the vertecies to black. The points may be moved with the mouse, or modified with the key commands listed below. Current data are always available in bry.x, bry.y and bry.verts. Key commands: enter : switch to grid editing mode t : toggle visibility of verticies d : delete a vertex i : insert a vertex at a point on the polygon line Methods: p.dump(bry_file='bry.pickle) Write the current boundary informtion (bry.verts) to a cPickle file bry_file. p.load(bry_file='bry.pickle) Read in boundary informtion (verts) from the cPickle file bry_file. Attributes: p.x, p.y : The x and y locations of the polygon verticies. p.verts : The verticies of the polygon (equiv. to zip(x, y)) p.area : the area of the polygon. Positive for a counterclockwise path, negative for a clockwise path. p.centroid : the (x, y) location of the polygon centroid. """ _showverts = True _epsilon = 5 # max pixel distance to count as a vertex hit def _init_boundary_interactor(self): """Send polyclick thread to boundary interactor""" # Get rid old mpl connections. self._ax.figure.canvas.mpl_disconnect(self._key_id) self._ax.figure.canvas.mpl_disconnect(self._click_id) # Shade the selected region in a polygon self._poly = Polygon(self.verts, alpha=0.2, fc='k', animated=True) self._ax.add_patch(self._poly) # modify the line properties self._line.set_markerfacecolor('k') # get the canvas and connect the callback events cid = self._poly.add_callback(self._poly_changed) self._ind = None # the active vert self._canvas.mpl_connect('draw_event', self._draw_callback) self._canvas.mpl_connect('button_press_event', self._button_press_callback) self._canvas.mpl_connect('key_press_event', self._key_press_callback) self._canvas.mpl_connect('button_release_event', self._button_release_callback) self._canvas.mpl_connect('motion_notify_event', self._motion_notify_callback) def _draw_callback(self, event): self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self._line.get_visible() Artist.update_from(self._line, poly) self._line.set_visible(vis) # don't use the poly visibility state def _get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' try: x, y = zip(*self._poly.xy) # display coords xt, yt = self._poly.get_transform().numerix_x_y(x, y) d = sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = nonzero(equal(d, amin(d))) ind = indseq[0] if d[ind] >= self._epsilon: ind = None return ind except: # display coords xy = asarray(self._poly.xy) xyt = self._poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = nonzero(equal(d, amin(d)))[0] ind = indseq[0] if d[ind] >= self._epsilon: ind = None return ind def _button_press_callback(self, event): 'whenever a mouse button is pressed' # if not self._showverts: return if event.inaxes == None: return if event.button != 1: return self._ind = self._get_ind_under_point(event) def _button_release_callback(self, event): 'whenever a mouse button is released' # if not self._showverts: return if event.button != 1: return self._ind = None def _key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self._showverts = not self._showverts elif event.key == 'd': ind = self._get_ind_under_point(event) if ind is not None: self._poly.xy = [ tup for i, tup in enumerate(self._poly.xy) if i != ind ] self._line.set_data(zip(*self._poly.xy)) elif event.key == 'i': xys = self._poly.get_transform().seq_xy_tups(self._poly.xy) p = event.x, event.y # display coords for i in range(len(xys) - 1): s0 = xys[i] s1 = xys[i + 1] d = dist_point_to_segment(p, s0, s1) if d <= self._epsilon: self._poly.xy.insert(i + 1, (event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) break s0 = xys[-1] s1 = xys[0] d = dist_point_to_segment(p, s0, s1) if d <= self._epsilon: self._poly.xy.append((event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) self._draw_callback(event) self._canvas.draw() def _motion_notify_callback(self, event): 'on mouse movement' if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x, y = event.xdata, event.ydata self._poly.xy[self._ind] = x, y x, y = zip(*self._poly.xy) self._line.set_data(x, y) self._canvas.restore_region(self._background) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _on_return(self, event): if event.key in ('enter', None): self._init_boundary_interactor() def _on_click(self, event): self.x.append(event.xdata) self.y.append(event.ydata) self._line.set_data(self.x, self.y) self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def __init__(self, *args, **kwargs): ax = kwargs.pop('ax', None) if ax is None: ax = pl.gca() self._ax = ax assert len( args ) <= 2, 'Input is either verticies (1 arg), or x and y (2 args)' if args is not (): if isinstance(args[0], str): verts = load(x) x, y = zip(*verts) x = list(x) y = list(y) elif len(args) == 1: # assume verticies input verts = args[0] x, y = zip(*verts) x = list(x) y = list(y) else: # x and y inputs x = list(args[0]) y = list(args[1]) assert len(x) == len(y), 'x and y must have the same length.' else: # no input x = [] y = [] self._line = pl.Line2D(x, y, marker='o', markerfacecolor='orange', animated=True) self._ax.add_line(self._line) self._canvas = self._line.figure.canvas self._key_id = self._ax.figure.canvas.mpl_connect( 'key_press_event', self._on_return) self._click_id = self._ax.figure.canvas.mpl_connect( 'button_press_event', self._on_click) if len(x) > 0: self._init_boundary_interactor() def dump(self, poly_file): f = open(bry_file, 'wb') cPickle.dump(self.verts, f, protocol=-1) f.close() def load(self, poly_file): verts = load(bry_file) x, y = zip(*verts) self._line.set_data(x, y) if hasattr(self, '_poly'): self._poly.xy = verts self._draw_callback(None) self._canvas.draw() def inside(self, *args): """docstring for inside""" assert len(args) > 0, 'must provide either verticies or x and y.' if len(args) == 1: # assume verticies input verts = args[0] else: # x and y inputs x = list(args[0]) y = list(args[1]) assert len(x) == len(y), 'x and y must have the same length.' verts = zip(x, y) return self._geom.inside(verts) def get_verts(self): return zip(self.x, self.y) verts = property(get_verts) def get_xdata(self): return self._line.get_xdata() x = property(get_xdata) def get_ydata(self): return self._line.get_ydata() y = property(get_ydata) def _get_geom(self): return Polygeom(self.verts) _geom = property(_get_geom) def _get_area(self): return self._geom.area area = property(_get_area) def _get_centroid(self): return self._geom.centroid centroid = property(_get_centroid)
class MaskCreator(object): """An interactive polygon editor. Parameters ---------- poly_xy : list of (float, float) List of (x, y) coordinates used as vertices of the polygon. max_ds : float Max pixel distance to count as a vertex hit. Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices """ def __init__(self, ax, poly_xy=None, max_ds=10): self.showverts = True self.max_ds = max_ds if poly_xy is None: poly_xy = default_vertices(ax) self.poly = Polygon(poly_xy, animated=True, fc='y', ec='none', alpha=0.4) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title("Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "Close figure when done.") self.ax = ax x, y = zip(*self.poly.xy) self.line = plt.Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.2, animated=True) self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def get_mask(self, shape): """Return image mask given by mask creator""" h, w = shape y, x = np.mgrid[:h, :w] points = np.transpose((x.ravel(), y.ravel())) mask = path.Path.contains_point(points) return mask.reshape(h, w) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() #Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def button_press_callback(self, event): 'whenever a mouse button is pressed' ignore = not self.showverts or event.inaxes is None or event.button != 1 if ignore: return self._ind = self.get_ind_under_cursor(event) def button_release_callback(self, event): 'whenever a mouse button is released' ignore = not self.showverts or event.button != 1 if ignore: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='t': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key=='d': ind = self.get_ind_under_cursor(event) if ind is None: return if ind == 0 or ind == self.last_vert_ind: print "Cannot delete root node" return self.poly.xy = [tup for i,tup in enumerate(self.poly.xy) if i!=ind] self._update_line() elif event.key=='i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # cursor coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d <= self.max_ds: self.poly.xy = np.array( list(self.poly.xy[:i+1]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i+1:])) self._update_line() break self.canvas.draw() def motion_notify_callback(self, event): 'on mouse movement' ignore = (not self.showverts or event.inaxes is None or event.button != 1 or self._ind is None) if ignore: return x,y = event.xdata, event.ydata if self._ind == 0 or self._ind == self.last_vert_ind: self.poly.xy[0] = x,y self.poly.xy[self.last_vert_ind] = x,y else: self.poly.xy[self._ind] = x,y self._update_line() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def _update_line(self): # save verts because polygon gets deleted when figure is closed self.verts = self.poly.xy self.last_vert_ind = len(self.poly.xy) - 1 self.line.set_data(zip(*self.poly.xy)) def get_ind_under_cursor(self, event): 'get the index of the vertex under cursor if within max_ds tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.max_ds: ind = None return ind
class PolygonEditor(object): ''' This edits the polygons drawn on the map ''' show_verts = True epsilon = 3 #threshold def __init__(self, axis, canvas): ''' initialises the editable polygon object ''' self.axis = axis self.polygon = None self.line = None self.xy_values = np.array([]) self._ind = None self.background = None #background copying self._callback_ids = list() self._callback_ids.append( canvas.mpl_connect('draw_event', self.draw_callback)) self._callback_ids.append( canvas.mpl_connect('button_press_event', self.button_press_callback)) self._callback_ids.append( canvas.mpl_connect('button_release_event', self.button_release_callback)) self._callback_ids.append( canvas.mpl_connect('motion_notify_event', self.motion_notify_callback)) self.canvas = canvas def add_point(self, xval, yval): """ adds an new point to the selection list and redraws the selection tool""" if self.xy_values.shape[0] == 0: self.xy_values = np.array([ (xval, yval), ]) else: self.xy_values = np.concatenate((self.xy_values, [ [xval, yval], ]), axis=0) self.refresh() def refresh(self): """ This method looks at the list of points available and depending on the number of points in the list creates a point or line or a polygon and draws them""" if self.xy_values.shape[0] == 0: # No points available self.reset_line() self.reset_polygon() elif self.xy_values.shape[0] <= 2: # point or line for 1 or 2 points xval, yval = zip(*self.xy_values) if self.line == None: self.line = Line2D(xval, yval, marker='o', markerfacecolor='r', animated=True) self.axis.add_line(self.line) else: self.line.set_data(zip(*self.xy_values)) self.reset_polygon() else: # more than 2 points if polygon is not created then creates one and draws if self.polygon == None: self.polygon = Polygon(self.xy_values, animated=True, alpha=polygon_alpha) self.polygon.add_callback(self.polygon_changed) self.axis.add_patch(self.polygon) else: self.polygon.xy = self.xy_values self.line.set_data(zip(*self.xy_values)) self.draw_callback(None) self.canvas.draw() def reset_line(self): """ resets the line i.e removes the line from the axes and resets to None """ if self.line != None: self.line.remove() self.line = None def reset_polygon(self): """ resets the polygon ie. removes the polygon from the axis and reset to None """ if self.polygon != None: self.polygon.remove() self.polygon = None def draw_line(self): """ draws the line if available """ if self.line != None: self.axis.draw_artist(self.line) def draw_polygon(self): """ draws polygon if available""" if self.polygon != None: self.axis.draw_artist(self.polygon) def disable(self): """ disables the events and the selection """ self.set_visibility(False) for callback_id in self._callback_ids: self.canvas.mpl_disconnect(callback_id) self.canvas = None def enable(self): """ enables the selection """ self.set_visibility(True) def set_visibility(self, status): """ sets the visibility of the selection object """ if self.polygon != None: self.polygon.set_visible(status) if self.line != None: self.line.set_visible(status) self.canvas.blit(self.axis.bbox) def draw_callback(self, dummy_event): """ this method draws the selection object """ self.background = self.canvas.copy_from_bbox(self.axis.bbox) self.draw_polygon() self.draw_line() self.canvas.blit(self.axis.bbox) def polygon_changed(self, poly): """ redraws the polygon """ vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) def get_index_under_point(self, event): """ gets the index of the point under the event (mouse click) """ if self.xy_values.shape[0] == 0: return None xy_values = self.xy_values xt_values, yt_values = xy_values[:, 0], xy_values[:, 1] dist = np.sqrt((xt_values - event.xdata)**2 + (yt_values - event.ydata)**2) indseq = np.nonzero(np.equal(dist, np.amin(dist)))[0] ind = indseq[0] if dist[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): """ callback to mouse press event """ if not self.show_verts: return if event.inaxes == None: return if event.button != 1: return self._ind = self.get_index_under_point(event) if self._ind == None: self.insert_datapoint(event) def button_release_callback(self, event): """ callback to mouse release event """ if not self.show_verts: return if event.button == 2: self.insert_datapoint(event) return if event.button == 3: self.delete_datapoint(event) return if event.button != 1: return self._ind = None def insert_datapoint(self, event): """ inserts a new data point between the segment that is closest in polygon """ if self.xy_values.shape[0] <= 2: self.add_point(event.xdata, event.ydata) else: event_point = event.xdata, event.ydata prev_d = dist_point_to_segment(event_point, self.xy_values[0], self.xy_values[-1]) prev_i = len(self.xy_values) for i in range(len(self.xy_values) - 1): seg_start = self.xy_values[i] seg_end = self.xy_values[i + 1] dist_p_s = dist_point_to_segment(event_point, seg_start, seg_end) if dist_p_s < prev_d: prev_i = i prev_d = dist_p_s self.xy_values = np.array( list(self.xy_values[:prev_i + 1]) + [(event.xdata, event.ydata)] + list(self.xy_values[prev_i + 1:])) self.refresh() def delete_datapoint(self, event): """ deletes the data point under the point in event """ ind = self.get_index_under_point(event) if ind is not None: self.xy_values = np.array( [tup for i, tup in enumerate(self.xy_values) if i != ind]) self.refresh() self.canvas.draw() def motion_notify_callback(self, event): """ callback for the mouse motion with button press. this is to move the edge points of the polygon""" if not self.show_verts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return xval, yval = event.xdata, event.ydata self.xy_values[self._ind] = xval, yval self.refresh() self.canvas.restore_region(self.background) self.axis.draw_artist(self.polygon) self.axis.draw_artist(self.line) self.canvas.blit(self.axis.bbox) def reset(self): """ resets the points in the selection deleting the line and polygon""" self.xy_values = np.array([]) self.reset_line() self.reset_polygon()
class MaskCreator(object): """An interactive polygon editor. Parameters ---------- poly_xy : list of (float, float) List of (x, y) coordinates used as vertices of the polygon. max_ds : float Max pixel distance to count as a vertex hit. Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices """ def __init__(self, ax, poly_xy=None, max_ds=10, line_width=2, line_color=(0, 0, 1), face_color=(1, .5, 0)): self.showverts = True self.max_ds = max_ds if poly_xy is None: poly_xy = default_vertices(ax) self.poly = Polygon(poly_xy, animated=True, fc=face_color, ec='none', alpha=0.4) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title("Click and drag a point to move it; or click once, then click again." "\n" "Close figure when done.") self.ax = ax x, y = zip(*self.poly.xy) #line_color = 'none' color = np.array(line_color) * .6 marker_face_color = line_color line_kwargs = {'lw': line_width, 'color': color, 'mfc': marker_face_color} self.line = plt.Line2D(x, y, marker='o', alpha=0.8, animated=True, **line_kwargs) self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) #canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def get_mask(self, shape): """Return image mask given by mask creator""" mask = verts_to_mask(shape, self.verts) return mask def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() #Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def draw_callback(self, event): #print('[mask] draw_callback(event=%r)' % event) self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self._ind is None: self._ind = None; return ignore = not self.showverts or event.inaxes is None or event.button != 1 if ignore: return self._ind = self.get_ind_under_cursor(event) if self._ind != None: self.indX, self.indY = self.poly.xy[self._ind] self.mouseX,self.mouseY = event.xdata, event.ydata def button_release_callback(self, event): 'whenever a mouse button is released' ignore = not self.showverts or event.button != 1 if ignore: return if self._ind is None: return currX, currY = self.poly.xy[self._ind] #print (currX, ' ', currY) #print (math.fabs(self.indX - currX)<3, ' and ', math.fabs(self.indY-currY)<3) if math.fabs(self.indX - currX)<3 and math.fabs(self.indY-currY)<3: return self._ind = None # def key_press_callback(self, event): # 'whenever a key is pressed' # if not event.inaxes: # return # if event.key == 't': # self.showverts = not self.showverts # self.line.set_visible(self.showverts) # if not self.showverts: # self._ind = None # elif event.key == 'd': # ind = self.get_ind_under_cursor(event) # if ind is None: # return # if ind == 0 or ind == self.last_vert_ind: # print('[mask] Cannot delete root node') # return # self.poly.xy = [tup for i, tup in enumerate(self.poly.xy) if i != ind] # self._update_line() # elif event.key == 'i': # xys = self.poly.get_transform().transform(self.poly.xy) # p = event.x, event.y # cursor coords # for i in range(len(xys) - 1): # s0 = xys[i] # s1 = xys[i + 1] # d = dist_point_to_segment(p, s0, s1) # if d <= self.max_ds: # self.poly.xy = np.array( # list(self.poly.xy[:i + 1]) + # [(event.xdata, event.ydata)] + # list(self.poly.xy[i + 1:])) # self._update_line() # break # self.canvas.draw() def motion_notify_callback(self, event): 'on mouse movement' # ignore = (not self.showverts or event.inaxes is None or # event.button != 1 or self._ind is None) ignore = (not self.showverts or event.inaxes is None) if ignore: #print ('verts ', not self.showverts, ' inaxes ', event.inaxes, ' event.buton ' ,event.button !=1, ' ind ', self._ind ) return if self._ind is None and event.button ==1: 'move all vertices' self.move_rectangle(event) self._update_line() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) self._ind = None 'set new mouse loc' self.mouseX,self.mouseY = event.xdata, event.ydata # if self._ind is None: # #create new poly # poly2_xy = vertices_under_cursor(event) # self.poly2 = Polygon(poly2_xy, animated=True, # fc=face_color, ec='none', alpha=0.4) # #ax.add_patch(self.poly2) # #Grab 3rd(?) ind if self._ind is None: return self.calculate_move(event) self._update_line() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def move_rectangle(self,event): selectedX, selectedY = (self.poly.xy[1]) beforeX, beforeY = (self.poly.xy[0]) afterX, afterY = (self.poly.xy[2]) acrossX, acrossY = (self.poly.xy[3]) listX = [selectedX, beforeX, afterX, acrossX] listY = [selectedY, beforeY, afterY, acrossY] maxX = max(listX) maxY = max(listY) minX = min(listX) minY = min(listY) x, y = event.xdata, event.ydata if x < minX or x> maxX or y<minY or y>maxY: return # Change selected self.poly.xy[1] = selectedX+(x-self.mouseX), selectedY+(y-self.mouseY) # Change before vert self.poly.xy[0] = beforeX+(x-self.mouseX), beforeY+(y-self.mouseY) self.poly.xy[self.last_vert_ind] = beforeX+(x-self.mouseX), beforeY+(y-self.mouseY) # Change after vert self.poly.xy[2] = afterX+(x-self.mouseX), afterY+(y-self.mouseY) #Change across vert self.poly.xy[3] = acrossX+(x-self.mouseX), acrossY+(y-self.mouseY) def calculate_move(self,event): indBefore = self._ind-1 if(indBefore < 0): indBefore = len(self.poly.xy)-2 indAfter = (self._ind+1)%4 selectedX, selectedY = (self.poly.xy[self._ind]) beforeX, beforeY = (self.poly.xy[indBefore]) afterX, afterY = (self.poly.xy[indAfter]) changeBefore = -1 keepX, changeY = -1, -1 changeAfter = -1 changeX, keepY = -1, -1 if beforeX != selectedX: changeBefore = indBefore keepX, changeY = self.poly.xy[indBefore] changeAfter = indAfter changeX, keepY = self.poly.xy[indAfter] else: changeBefore = indAfter keepX, changeY = self.poly.xy[indAfter] changeAfter = indBefore changeX, keepY = self.poly.xy[indBefore] x, y = event.xdata, event.ydata # Change selected if self._ind == 0 or self._ind == self.last_vert_ind: self.poly.xy[0] = x, y self.poly.xy[self.last_vert_ind] = x, y else: self.poly.xy[self._ind] = x, y # Change vert if changeBefore == 0 or changeBefore == self.last_vert_ind: self.poly.xy[0] = keepX, y self.poly.xy[self.last_vert_ind] = keepX, y else: self.poly.xy[changeBefore] = keepX, y # Change horiz if changeAfter == 0 or changeAfter == self.last_vert_ind: self.poly.xy[0] = x, keepY self.poly.xy[self.last_vert_ind] = x, keepY else: self.poly.xy[changeAfter] = x, keepY def _update_line(self): # save verts because polygon gets deleted when figure is closed self.verts = self.poly.xy self.last_vert_ind = len(self.poly.xy) - 1 self.line.set_data(zip(*self.poly.xy)) def get_ind_under_cursor(self, event): 'get the index of the vertex under cursor if within max_ds tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt - event.x) ** 2 + (yt - event.y) ** 2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.max_ds: ind = None return ind
class PolyRegionEditor(uiutils.UI): def __init__(self, img, prj=None, max_ds=10, current_folder=None): self.max_ds = max_ds self.showverts = True self.current_folder = current_folder uiutils.UI.__init__(self, 750, 600, "PolyRegion Editor") vbox = QtGui.QVBoxLayout() self.setLayout(vbox) canva_box = QtGui.QHBoxLayout() vbox.addLayout(canva_box) self.canvas = plotutils.BaseCustomCanvas() canva_box.addWidget(self.canvas) ctl = QtGui.QHBoxLayout() vbox.addLayout(ctl) ctl.addWidget(plotutils.NavigationToolbar(self.canvas, self)) self.title_entry = uiutils.EntryDescription("Title") ctl.addWidget(self.title_entry) self.color_entry = uiutils.EntryDescription("Color") ctl.addWidget(self.color_entry) save_bn = QtGui.QPushButton("Save") save_bn.clicked.connect(self.on_save_clicked) ctl.addWidget(save_bn) load_bn = QtGui.QPushButton("Load") load_bn.clicked.connect(self.on_load_clicked) ctl.addWidget(load_bn) save_bn = QtGui.QPushButton("New") save_bn.clicked.connect(self.on_new_clicked) ctl.addWidget(save_bn) self.ax = self.canvas.figure.subplots() if prj is None: prj = img.get_projection() self.prj = prj self.img = img plotutils.imshow_image(self.ax, self.img, projection=self.prj, title=False) self.poly = Polygon([[0, 0]], animated=True, fc='b', ec='none', alpha=0.4) self.ax.add_patch(self.poly) self.ax.set_clip_on(False) self.ax.set_title("Click and drag a point to move it; " "'i' to insert; 'd' to delete.") x, y = zip(*self.poly.xy) self.line = plt.Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.8, animated=True, lw=2, markersize=self.max_ds) self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas self.load_default() self.canvas.setFocus() def get_axes(self): return self.ax def load_poly_region(self, poly_region): self.title_entry.set_text(poly_region.title) self.color_entry.set_text(poly_region.color) self.poly.set_fc(poly_region.color) vertices = list(poly_region.vertices) vertices.append(vertices[0]) self.poly.xy = vertices self._update_line() self.canvas.draw() def load_default(self): default_poly_region = imgutils.PolyRegion.default_from_ax(self.ax) self.load_poly_region(default_poly_region) def on_load_clicked(self, bn): filename = uiutils.open_file(parent=self, current_folder=self.current_folder) if filename is not None: try: poly_region = imgutils.PolyRegion.from_file(filename, self.img.get_coordinate_system()) self.load_poly_region(poly_region) except Exception, e: msg = "Failed to load region %s\n%s" % (filename, e) print msg uiutils.error_msg(msg, self)
class BsplineInteractor(object): """ An polygon editor. Key-bindings 't' toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' delete the vertex under point 'i' insert a vertex at point. You must be within epsilon of the line connecting two existing vertices """ #showverts = True #epsilon = 5 def __init__(self, Lspline, mouse_update_listener=None, optimization_listener=None): self.Lspline = Lspline self.curve = Lspline.curve self.mouse_update_listener = [mouse_update_listener] self.optimization_listener = [optimization_listener] self.showverts = True self.epsilon = 5 self.verbose = True self.fig = plt.figure() #figsize=(4, 4) #self.fig2 = plt.figure() self.ax = self.fig.add_subplot(111) self.fig.add_axes(self.ax) self.controlpoly = Polygon(list( zip(self.curve.vertices[:, 0], self.curve.vertices[:, 1])), animated=True, fill=False, closed=False) self.curve_r = Polygon(list(zip(self.curve.r[:, 0], self.curve.r[:, 1])), animated=True, fill=False, closed=False) self.ax.add_patch(self.controlpoly) self.ax.add_patch(self.curve_r) self.line_poly = Line2D(self.curve.vertices[:, 0], self.curve.vertices[:, 1], color='blue', marker='o', markerfacecolor='r', alpha=.5, animated=True) self.line_r = Line2D(self.curve.r[:, 0], self.curve.r[:, 1], color='black', alpha=.75, markerfacecolor='r', animated=True) self.ax.add_line(self.line_poly) self.ax.add_line(self.line_r) self.cid = self.controlpoly.add_callback(self.poly_changed) self.did = self.curve_r.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.controlpoly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.setconstraints = plt.axes([0.50, 0.05, 0.1, 0.075]) self.sconstraints = Button(self.setconstraints, 'SET') self.sconstraints.on_clicked(self.set_contraints) self.axconstraints = plt.axes([0.61, 0.05, 0.1, 0.075]) self.bconstraints = Button(self.axconstraints, 'Form Parm') self.bconstraints.on_clicked(self.add_constraints) self.axprev = plt.axes([0.72, 0.05, 0.1, 0.075]) self.bprev = Button(self.axprev, 'Optimize') self.bprev.on_clicked(self.do_optimization) self.canvas = canvas self.ax.set_title('Lspline Plot') xmin, xmax, ymin, ymax = self.Lspline.extreme_C0() self.ax.set_xlim((xmin - 2., xmax + 2.)) self.ax.set_ylim((ymin - 2., ymax + 2.)) return def add_constraints(self, event): if self.verbose == True: print 'adding constraints' self.FPmaker = FormParameterInteractor(self.Lspline) return def set_contraints(self, event): if self.verbose == True: print 'setting constraints' self.Lspline = self.FPmaker() return def do_optimization(self, event): if self.verbose == True: print 'doing optimization' self.Lspline.optimize() self.compute_curve_plot() if self.optimization_listener[0] is not None: for func in self.optimization_listener: func(self.Lspline.curve.vertices) return def plot_contraints(self, event): if self.verbose == True: print 'plotting constraints' self.Lspline.display(mytitle='Test 1') return def compute_curve_plot(self): self.controlpoly.xy = self.curve.vertices self.line_poly.set_data(zip(*self.curve.vertices)) self.canvas.restore_region(self.background) #self.curve.vertices = self.controlpoly.xy self.curve.allCurveArray() self.curve_r = Polygon(self.curve.r, animated=True, fill=False, closed=False) self.curve_r.yx = self.curve.r self.line_r.set_data(self.curve.r[:, 0], self.curve.r[:, 1]) self.ax.draw_artist(self.controlpoly) self.ax.draw_artist(self.line_poly) self.ax.draw_artist(self.curve_r) self.ax.draw_artist(self.line_r) self.canvas.blit(self.ax.bbox) return def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.controlpoly) self.ax.draw_artist(self.line_poly) self.ax.draw_artist(self.curve_r) self.ax.draw_artist(self.line_r) #self.canvas.blit(self.ax.bbox) return def poly_changed(self, controlpoly, curve_r): vis = self.line_poly.get_visible() Artist.update_from(self.line_poly, controlpoly) self.line_poly.set_visible(vis) vis = self.line_r.get_visible() Artist.update_from(self.line_r, curve_r) self.line_r.set_visible(vis) return def get_ind_under_point(self, event): """get the index of the vertex under point if within epsilon tolerance """ xy = np.asarray(self.controlpoly.xy) xyt = self.controlpoly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] print xt, yt print event.x, event.y d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if self.verbose: print 'button_press_callback GUISPLINE' print not self.showverts print event.inaxes == None print event.button != 1 if not self.showverts: return if event.inaxes == None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is reileased' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line_poly.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_point(event) if ind is not None: self.controlpoly.xy = [ tup for i, tup in enumerate(self.controlpoly.xy) if i != ind ] self.line_poly.set_data(zip(*self.controlpoly.xy)) elif event.key == 'i': xys = self.controlpoly.get_transform().transform( self.controlpoly.xy) p = event.x, event.y # display coords for i in range(len(xys)): s0 = xys[i - 1] s1 = xys[i] d = dist_point_to_segment(p, s0, s1) if d <= self.epsilon: self.controlpoly.xy = np.array( list(self.controlpoly.xy[:i]) + [(event.xdata, event.ydata)] + list(self.controlpoly.xy[i:])) self.line_poly.set_data(zip(*self.controlpoly.xy)) break self.canvas.draw() def motion_notify_callback(self, event): 'on mouse movement' if self.verbose: print 'motion_notify_callback' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x, y = event.xdata, event.ydata self.curve.vertices[self._ind, 0] = x self.curve.vertices[self._ind, 1] = y self.compute_curve_plot() if self.mouse_update_listener[0] is not None: for func in self.mouse_update_listener: func(x, y, self._ind) return
class PolygonInteractor(QtCore.QObject): """ Polygon Interactor Parameters ---------- axtmp : matplotlib axis matplotlib axis pntxy : """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit polyi_changed = QtCore.pyqtSignal(list) def __init__(self, axtmp, pntxy): QtCore.QObject.__init__(self) self.ax = axtmp self.poly = Polygon([(1, 1)], animated=True) self.ax.add_patch(self.poly) self.canvas = self.poly.figure.canvas self.poly.set_alpha(0.5) self.pntxy = pntxy self.ishist = True self.background = self.canvas.copy_from_bbox(self.ax.bbox) xtmp, ytmp = list(zip(*self.poly.xy)) self.line = Line2D(xtmp, ytmp, marker='o', markerfacecolor='r', color='y', animated=True) self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert self.canvas.mpl_connect('button_press_event', self.button_press_callback) self.canvas.mpl_connect('button_release_event', self.button_release_callback) self.canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) def draw_callback(self): """ Draw callback """ self.background = self.canvas.copy_from_bbox(self.ax.bbox) QtWidgets.QApplication.processEvents() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() def new_poly(self, npoly): """ New Polygon """ self.poly.set_xy(npoly) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.draw() self.update_plots() def poly_changed(self, poly): """ Changed Polygon """ # this method is called whenever the polygon object is called # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): """get the index of vertex under point if within epsilon tolerance""" # display coords xytmp = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xytmp) xtt, ytt = xyt[:, 0], xyt[:, 1] dtt = np.sqrt((xtt - event.x) ** 2 + (ytt - event.y) ** 2) indseq = np.nonzero(np.equal(dtt, np.amin(dtt)))[0] ind = indseq[0] if dtt[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): """whenever a mouse button is pressed""" if event.inaxes is None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) if self._ind is None: xys = self.poly.get_transform().transform(self.poly.xy) ptmp = event.x, event.y # display coords if len(xys) == 1: self.poly.xy = np.array( [(event.xdata, event.ydata)] + [(event.xdata, event.ydata)]) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() return dmin = -1 imin = -1 for i in range(len(xys) - 1): s0tmp = xys[i] s1tmp = xys[i + 1] dtmp = dist_point_to_segment(ptmp, s0tmp, s1tmp) if dmin == -1: dmin = dtmp imin = i elif dtmp < dmin: dmin = dtmp imin = i i = imin self.poly.xy = np.array(list(self.poly.xy[:i + 1]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i + 1:])) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() def button_release_callback(self, event): """Whenever a mouse button is released""" if event.button != 1: return self._ind = None self.update_plots() def update_plots(self): """ Update Plots """ polymask = Path(self.poly.xy).contains_points(self.pntxy) self.polyi_changed.emit(polymask.tolist()) def motion_notify_callback(self, event): """on mouse movement""" if self._ind is None: return if event.inaxes is None: return if event.button != 1: return xtmp, ytmp = event.xdata, event.ydata self.poly.xy[self._ind] = xtmp, ytmp if self._ind == 0: self.poly.xy[-1] = xtmp, ytmp self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update()
class MaskDrawer(object): """An interactive polygon mask drawer on an image. Parameters ---------- ax : matplotlib plot axis Inpimg : 2d numpy array Input image to overlay for drawing mask Mask : Boolean numpy array same size of inpimg A Mask which will be used as initial mask and updated upon confirmation max_ds : float Max pixel distance to count as a vertex hit. PolyAtStart : List of vertices A list of square vertices to draw the initial polygon Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices 'n' : Invert the region selected by polynomial for masking 'c' : Confirm the polygon and update the mask """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit def __init__(self, ax, Inpimg, Mask, max_ds=10, PolyAtStart=[(50, 50), (100, 50), (100, 100), (50, 100)]): self.showverts = True self.max_ds = max_ds self.Mask = Mask self.img = Inpimg self.maskinvert = False # imshow the image self.imgplot = ax.imshow(np.ma.filled( np.ma.array(self.img, mask=self.Mask, fill_value=np.nan)), cmap=cm.gray) self.poly = Polygon(PolyAtStart, animated=True, fc='y', ec='none', alpha=0.5) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title( "Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "'n' to invert the region for masking, 'c' to confirm & apply the mask." ) self.ax = ax x, y = zip(*self.poly.xy) self.line = Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.7, animated=True) # self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes == None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_point(event) if ind is not None: self.poly.xy = [ tup for i, tup in enumerate(self.poly.xy) if i != ind ] self.line.set_data(zip(*self.poly.xy)) elif event.key == 'i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # display coords for i in range(len(xys) - 1): s0 = xys[i] s1 = xys[i + 1] d = dist_point_to_segment(p, s0, s1) if d <= self.epsilon: self.poly.xy = np.array( list(self.poly.xy[:i]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i:])) self.line.set_data(zip(*self.poly.xy)) break elif event.key == 'n': """ Flips the region inside out of the polygon to be masked """ print('Inverting the mask region') self.maskinvert = not self.maskinvert elif event.key == 'c': """ Confirm the drawn polynomial shape and add update the mask """ self.UpdateMask() #Update the imshowed image with new mask self.imgplot.set_data( np.ma.filled( np.ma.array(self.img, mask=self.Mask, fill_value=np.nan))) self.imgplot.figure.canvas.draw() self.canvas.draw() def UpdateMask(self): """ Updates the maks array with points insied the polygon """ print('Updating the original Mask..') Path = self.poly.get_path() h, w = self.Mask.shape y, x = np.mgrid[:h, :w] XYpoints = np.transpose((x.ravel(), y.ravel())) NewMask = Path.contains_points(XYpoints) if self.maskinvert: NewMask = ~NewMask # Combine the two mask by taing an elemet wise or self.Mask = self.Mask | NewMask.reshape((h, w)) def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x, y = event.xdata, event.ydata self.poly.xy[self._ind] = x, y self.line.set_data(zip(*self.poly.xy)) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox)
class BoundaryClick(object): """ bry = BoundaryClick(x=[], y=[], beta=None, ax=gca(), **gridgen_options) If x, y and beta are not given, and interactive polygon creation session is started. The initial boundary is sketched out by clicking the points (in a counterclockwise manner, usually starting in the upper lefthand corner of the boundary). At this point the verticies are marked by orange circles. Switch to editing mode by hitting return. This changes the vertecies to black. At this point, beta values may be created to define the corners of the grid (see below). Current data are always available in bry.x, bry.y (combined in p.verts) and bry.beta attributes. Key commands: enter : switch to grid editing mode t : toggle visibility of verticies d : delete a vertex i : insert a vertex at a point on the polygon line p : define vertex as beta=1 (a Positive turn, marked with green triangle) m : define vertex as beta=1 (a Negative turn, marked with red triangle) z : define vertex as beta=0 (no corner, marked with green triangle) G : generate grid from the current boundary using gridgen T : toggle visability of the current grid R : remove the gridlines from the figure Methods: bry.dump(bry_file) Write the current boundary informtion (bry.x, bry.y, bry.beta) to a cPickle file bry_file. bry.load(bry_file) Read in boundary informtion (x, y, beta) from the cPickle file bry_file. bry.remove_grid() Remove gridlines from axes. """ _showverts = True _showbetas = True _showgrid = True _epsilon = 5 # max pixel distance to count as a vertex hit def _update_beta_lines(self): """Update the m/pline by finding the (x,y) points where self.beta== -/+ 1""" x = self._line.get_xdata() y = self._line.get_ydata() xp = [x[n] for n in range(len(x)) if self.beta[n]==1] yp = [y[n] for n in range(len(y)) if self.beta[n]==1] self._pline.set_data(xp, yp) xm = [x[n] for n in range(len(x)) if self.beta[n]==-1] ym = [y[n] for n in range(len(y)) if self.beta[n]==-1] self._mline.set_data(xm, ym) xz = [x[n] for n in range(len(x)) if self.beta[n]==0] yz = [y[n] for n in range(len(y)) if self.beta[n]==0] self._zline.set_data(xz, yz) if len(x)-1 < self.gridgen_options['ul_idx']: self.gridgen_options['ul_idx'] = len(x)-1 xs = x[self.gridgen_options['ul_idx']] ys = y[self.gridgen_options['ul_idx']] self._sline.set_data(xs, ys) def remove_grid(self): """Remove a generated grid from the BoundaryClick figure""" if hasattr(self, '_gridlines'): for line in self._gridlines: self._ax.lines.remove(line) delattr(self, '_gridlines') def _init_boundary_interactor(self): """Send polyclick thread to boundary interactor""" # Get rid old mpl connections. self._ax.figure.canvas.mpl_disconnect(self._key_id) self._ax.figure.canvas.mpl_disconnect(self._click_id) # Shade the selected region in a polygon self._poly = Polygon(self.verts, alpha=0.1, fc='k', animated=True) self._ax.add_patch(self._poly) # change line to animated # self._line.set_animated(True) # self._line.set_markerfacecolor('k') self._line.set_marker('None') # Link in the two lines that will show the beta values # pline for positive turns, mline for negative (minus) turns. self._pline = Line2D([], [], marker='^', ms=12, mfc='g', animated=True, lw=0) self._mline = Line2D([], [], marker='v', ms=12, mfc='r', animated=True, lw=0) self._zline = Line2D([], [], marker='o', mfc='k', animated=True, lw=0) self._sline = Line2D([], [], marker='s', mfc='k', animated=True, lw=0) self._update_beta_lines() self._ax.add_artist(self._pline) self._ax.add_artist(self._mline) self._ax.add_artist(self._zline) self._ax.add_artist(self._sline) # get the canvas and connect the callback events cid = self._poly.add_callback(self._poly_changed) self._ind = None # the active vert self._canvas.mpl_connect('draw_event', self._draw_callback) self._canvas.mpl_connect('button_press_event', self._button_press_callback) self._canvas.mpl_connect('key_press_event', self._key_press_callback) self._canvas.mpl_connect('button_release_event', self._button_release_callback) self._canvas.mpl_connect('motion_notify_event', self._motion_notify_callback) def _on_return(self, event): if event.key in ('enter', None): self._init_boundary_interactor() def _on_click(self, event): self.x.append(event.xdata) self.y.append(event.ydata) self.beta.append(0) self._line.set_data(self.x, self.y) self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _draw_callback(self, event): self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._pline) self._ax.draw_artist(self._mline) self._ax.draw_artist(self._zline) self._ax.draw_artist(self._sline) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self._line.get_visible() Artist.update_from(self._line, poly) self._line.set_visible(vis) # don't use the poly visibility state def _get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' try: x, y = zip(*self._poly.xy) # display coords xt, yt = self._poly.get_transform().numerix_x_y(x, y) d = sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = nonzero(equal(d, amin(d))) ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind except: # display coords xy = asarray(self._poly.xy) xyt = self._poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = nonzero(equal(d, amin(d)))[0] ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind def _button_press_callback(self, event): 'whenever a mouse button is pressed' # if not self._showverts: return if event.inaxes==None: return if event.button != 1: return self._ind = self._get_ind_under_point(event) def _button_release_callback(self, event): 'whenever a mouse button is released' # if not self._showverts: return if event.button != 1: return self._ind = None def _key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='shift': return if event.key=='t': self._showbetas = not self._showbetas self._pline.set_visible(self._showbetas) self._mline.set_visible(self._showbetas) self._zline.set_visible(self._showbetas) self._sline.set_visible(self._showbetas) elif event.key=='d': ind = self._get_ind_under_point(event) if ind is not None: self._poly.xy = [tup for i,tup in enumerate(self._poly.xy) if i!=ind] self._line.set_data(zip(*self._poly.xy)) self.beta = [beta for i,beta in enumerate(self.beta) if i!=ind] elif event.key=='p': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = 1.0 elif event.key=='m': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = -1.0 elif event.key=='z': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = 0.0 elif event.key=='s': ind = self._get_ind_under_point(event) if ind is not None: self.gridgen_options['ul_idx'] = ind elif event.key=='i': xys = self._poly.get_transform().seq_xy_tups(self._poly.xy) p = event.x, event.y # display coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy.insert(i+1, (event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) self.beta.insert(i+1, 0) break s0 = xys[-1] s1 = xys[0] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy.append((event.xdata, event.ydata)) self._line.set_data(zip(*self._poly.xy)) self.beta.append(0) elif event.key=='G' or event.key == '1': options = copy.deepcopy(self.gridgen_options) shp = options.pop('shp') self.grd = pyroms.gridgen(self.x, self.y, self.beta, shp, **options) self.remove_grid() self._showgrid = True gridlineprops = {'linestyle':'-', 'color':'k', 'lw':0.2} self._gridlines = [] for line in self._ax._get_lines(*(self.grd.x, self.grd.y), **gridlineprops): self._ax.add_line(line) self._gridlines.append(line) for line in self._ax._get_lines(*(self.grd.x.T, self.grd.y.T), **gridlineprops): self._ax.add_line(line) self._gridlines.append(line) elif event.key=='R': self.remove_grid() elif event.key=='T' or event.key == '2': self._showgrid = not self._showgrid if hasattr(self, '_gridlines'): for line in self._gridlines: line.set_visible(self._showgrid) self._update_beta_lines() self._draw_callback(event) self._canvas.draw() def _motion_notify_callback(self, event): 'on mouse movement' # if not self._showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x,y = event.xdata, event.ydata self._poly.xy[self._ind] = x,y x, y = zip(*self._poly.xy) self._line.set_data(x, y) self._update_beta_lines() self._canvas.restore_region(self._background) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._pline) self._ax.draw_artist(self._mline) self._ax.draw_artist(self._zline) self._ax.draw_artist(self._sline) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def __init__(self, x=[], y=[], beta=None, ax=None, **gridgen_options): if isinstance(x, str): x, y, beta = load(x) if ax is None: ax = pl.gca() self._ax = ax # Set default gridgen option, and copy over specified options. self.gridgen_options = {'ul_idx': 0, 'shp': (32, 32), 'verbose': True} for key, value in gridgen_options.iteritems(): self.gridgen_options[key] = gridgen_options[key] x = list(x); y = list(y) assert len(x)==len(y), 'arrays must be equal length' if beta is None: self.beta = [0 for xi in x] else: assert len(x)==len(beta), 'beta must have same length as x and y' self.beta = list(beta) self._line = pl.Line2D(x, y, marker='o', markerfacecolor='orange', animated=True) self._ax.add_line(self._line) self._canvas = self._line.figure.canvas self._key_id = self._ax.figure.canvas.mpl_connect('key_press_event', self._on_return) self._click_id = self._ax.figure.canvas.mpl_connect('button_press_event', self._on_click) if len(x) > 0: self._init_boundary_interactor() def dump(self, bry_file): f = open(bry_file, 'wb') bry_dict = {'x': self.x, 'y': self.y, 'beta': self.beta} cPickle.dump(bry_dict, f, protocol=-1) f.close() def load(self, bry_file): bry_dict = load(bry_file) x = bry_dict['x'] y = bry_dict['y'] beta = bry_dict['beta'] self._line.set_data(x, y) if hasattr(self, '_poly'): self._poly.xy = zip(x, y) self._update_beta_lines() self._draw_callback(None) self._canvas.draw() def _get_verts(self): return zip(self.x, self.y) verts = property(_get_verts) def get_xdata(self): return self._line.get_xdata() x = property(get_xdata) def get_ydata(self): return self._line.get_ydata() y = property(get_ydata)
class PolygonInteractor(QtCore.QObject): """ Polygon Interactor Parameters ---------- axtmp : matplotlib axis matplotlib axis pntxy : """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit polyi_changed = QtCore.pyqtSignal(list) def __init__(self, axtmp, pntxy): QtCore.QObject.__init__(self) self.ax = axtmp self.poly = Polygon([(1, 1)], animated=True) self.ax.add_patch(self.poly) self.canvas = self.poly.figure.canvas self.poly.set_alpha(0.5) self.pntxy = pntxy self.ishist = True self.background = self.canvas.copy_from_bbox(self.ax.bbox) xtmp, ytmp = list(zip(*self.poly.xy)) self.line = Line2D(xtmp, ytmp, marker='o', markerfacecolor='r', color='y', animated=True) self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert self.canvas.mpl_connect('button_press_event', self.button_press_callback) self.canvas.mpl_connect('button_release_event', self.button_release_callback) self.canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) def draw_callback(self): """ Draw callback """ self.background = self.canvas.copy_from_bbox(self.ax.bbox) QtWidgets.QApplication.processEvents() self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() def new_poly(self, npoly): """ New Polygon """ self.poly.set_xy(npoly) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.draw() self.update_plots() def poly_changed(self, poly): """ Changed Polygon """ # this method is called whenever the polygon object is called # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): """get the index of vertex under point if within epsilon tolerance""" # display coords xytmp = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xytmp) xtt, ytt = xyt[:, 0], xyt[:, 1] dtt = np.sqrt((xtt - event.x) ** 2 + (ytt - event.y) ** 2) indseq = np.nonzero(np.equal(dtt, np.amin(dtt)))[0] ind = indseq[0] if dtt[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): """whenever a mouse button is pressed""" if event.inaxes is None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) if self._ind is None: xys = self.poly.get_transform().transform(self.poly.xy) ptmp = self.poly.get_transform().transform([event.xdata, event.ydata]) # ptmp = event.x, event.y # display coords if len(xys) == 1: self.poly.xy = np.array( [(event.xdata, event.ydata)] + [(event.xdata, event.ydata)]) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() return dmin = -1 imin = -1 for i in range(len(xys) - 1): s0tmp = xys[i] s1tmp = xys[i + 1] dtmp = dist_point_to_segment(ptmp, s0tmp, s1tmp) if dmin == -1: dmin = dtmp imin = i elif dtmp < dmin: dmin = dtmp imin = i i = imin # breakpoint() self.poly.xy = np.array(list(self.poly.xy[:i + 1]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i + 1:])) self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update() def button_release_callback(self, event): """Whenever a mouse button is released""" if event.button != 1: return self._ind = None self.update_plots() def update_plots(self): """ Update Plots """ polymask = Path(self.poly.xy).contains_points(self.pntxy) self.polyi_changed.emit(polymask.tolist()) def motion_notify_callback(self, event): """on mouse movement""" if self._ind is None: return if event.inaxes is None: return if event.button != 1: return xtmp, ytmp = event.xdata, event.ydata self.poly.xy[self._ind] = xtmp, ytmp if self._ind == 0: self.poly.xy[-1] = xtmp, ytmp self.line.set_data(list(zip(*self.poly.xy))) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.update()
class BoundaryInteractor(object): """ Interactive grid creation bry = BoundaryClick(x=[], y=[], beta=None, ax=gca(), **gridgen_options) The initial boundary polygon points (x and y) are counterclockwise, starting in the upper left corner of the boundary. Key commands: t : toggle visibility of verticies d : delete a vertex i : insert a vertex at a point on the polygon line p : set vertex as beta=1 (a Positive turn, marked with green triangle) m : set vertex as beta=1 (a Negative turn, marked with red triangle) z : set vertex as beta=0 (no corner, marked with a black dot) G : generate grid from the current boundary using gridgen T : toggle visability of the current grid Methods: bry.dump(bry_file) Write the current boundary informtion (bry.x, bry.y, bry.beta) to a cPickle file bry_file. bry.load(bry_file) Read in boundary informtion (x, y, beta) from the cPickle file bry_file. bry.remove_grid() Remove gridlines from axes. Attributes: bry.x : the X boundary points bry.y : the Y boundary points bry.verts : the verticies of the grid bry.grd : the CGrid object """ _showverts = True _showbetas = True _showgrid = True _epsilon = 5 # max pixel distance to count as a vertex hit def _update_beta_lines(self): """Update m/pline by finding the points where self.beta== -/+ 1""" x, y = zip(*self._poly.xy) num_points = len(x)-1 # the first and last point are repeated xp = [x[n] for n in range(num_points) if self.beta[n]==1] yp = [y[n] for n in range(num_points) if self.beta[n]==1] self._pline.set_data(xp, yp) xm = [x[n] for n in range(num_points) if self.beta[n]==-1] ym = [y[n] for n in range(num_points) if self.beta[n]==-1] self._mline.set_data(xm, ym) xz = [x[n] for n in range(num_points) if self.beta[n]==0] yz = [y[n] for n in range(num_points) if self.beta[n]==0] self._zline.set_data(xz, yz) if len(x)-1 < self.gridgen_options['ul_idx']: self.gridgen_options['ul_idx'] = len(x)-1 xs = x[self.gridgen_options['ul_idx']] ys = y[self.gridgen_options['ul_idx']] self._sline.set_data(xs, ys) def remove_grid(self): """Remove a generated grid from the BoundaryClick figure""" if hasattr(self, '_gridlines'): for line in self._gridlines: self._ax.lines.remove(line) delattr(self, '_gridlines') def _draw_callback(self, event): self._background = self._canvas.copy_from_bbox(self._ax.bbox) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._pline) self._ax.draw_artist(self._mline) self._ax.draw_artist(self._zline) self._ax.draw_artist(self._sline) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def _poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self._line.get_visible() Artist.update_from(self._line, poly) self._line.set_visible(vis) # don't use the poly visibility state def _get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' try: x, y = zip(*self._poly.xy) # display coords xt, yt = self._poly.get_transform().numerix_x_y(x, y) d = np.sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d))) ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind except: # display coords xy = np.asarray(self._poly.xy) xyt = self._poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind]>=self._epsilon: ind = None return ind def _button_press_callback(self, event): 'whenever a mouse button is pressed' # if not self._showverts: return if event.inaxes==None: return if event.button != 1: return self._ind = self._get_ind_under_point(event) def _button_release_callback(self, event): 'whenever a mouse button is released' # if not self._showverts: return if event.button != 1: return self._ind = None def _key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='shift': return if event.key=='t': self._showbetas = not self._showbetas self._line.set_visible(self._showbetas) self._pline.set_visible(self._showbetas) self._mline.set_visible(self._showbetas) self._zline.set_visible(self._showbetas) self._sline.set_visible(self._showbetas) elif event.key=='d': ind = self._get_ind_under_point(event) if ind is not None: self._poly.xy = [tup for i,tup in enumerate(self._poly.xy) \ if i!=ind] self._line.set_data(zip(*self._poly.xy)) self.beta = [beta for i,beta in enumerate(self.beta) \ if i!=ind] elif event.key=='p': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = 1.0 elif event.key=='m': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = -1.0 elif event.key=='z': ind = self._get_ind_under_point(event) if ind is not None: self.beta[ind] = 0.0 elif event.key=='s': ind = self._get_ind_under_point(event) if ind is not None: self.gridgen_options['ul_idx'] = ind elif event.key=='i': xys = self._poly.get_transform().transform(self._poly.xy) p = event.x, event.y # display coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy = np.array( list(self._poly.xy[:i+1]) + [(event.xdata, event.ydata)] + list(self._poly.xy[i+1:])) self._line.set_data(zip(*self._poly.xy)) self.beta.insert(i+1, 0) break s0 = xys[-1] s1 = xys[0] d = dist_point_to_segment(p, s0, s1) if d<=self._epsilon: self._poly.xy = np.array( list(self._poly.xy) + [(event.xdata, event.ydata)]) self._line.set_data(zip(*self._poly.xy)) self.beta.append(0) elif event.key=='G' or event.key == '1': options = deepcopy(self.gridgen_options) shp = options.pop('shp') if self.proj is None: x = self.x y = self.y self.grd = Gridgen(x, y, self.beta, shp, proj=self.proj, **options) else: lon, lat = self.proj(self.x, self.y, inverse=True) self.grd = Gridgen(lon, lat, self.beta, shp, proj=self.proj, **options) self.remove_grid() self._showgrid = True gridlineprops = {'linestyle':'-', 'color':'k', 'lw':0.2} self._gridlines = [] for line in self._ax._get_lines(*(self.grd.x, self.grd.y), **gridlineprops): self._ax.add_line(line) self._gridlines.append(line) for line in self._ax._get_lines(*(self.grd.x.T, self.grd.y.T), **gridlineprops): self._ax.add_line(line) self._gridlines.append(line) elif event.key=='T' or event.key == '2': self._showgrid = not self._showgrid if hasattr(self, '_gridlines'): for line in self._gridlines: line.set_visible(self._showgrid) self._update_beta_lines() self._draw_callback(event) self._canvas.draw() def _motion_notify_callback(self, event): 'on mouse movement' # if not self._showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x,y = event.xdata, event.ydata self._poly.xy[self._ind] = x, y if self._ind == 0: self._poly.xy[-1] = x, y x, y = zip(*self._poly.xy) self._line.set_data(x[:-1], y[:-1]) self._update_beta_lines() self._canvas.restore_region(self._background) self._ax.draw_artist(self._poly) self._ax.draw_artist(self._pline) self._ax.draw_artist(self._mline) self._ax.draw_artist(self._zline) self._ax.draw_artist(self._sline) self._ax.draw_artist(self._line) self._canvas.blit(self._ax.bbox) def __init__(self, x, y=None, beta=None, ax=None, proj=None, **gridgen_options): if isinstance(x, str): bry_dict = np.load(x) x = bry_dict['x'] y = bry_dict['y'] beta = bry_dict['beta'] assert len(x) >= 4, 'Boundary must have at least four points.' if ax is None: ax = plt.gca() self._ax = ax self.proj = proj # Set default gridgen option, and copy over specified options. self.gridgen_options = {'ul_idx': 0, 'shp': (32, 32)} for key, value in gridgen_options.iteritems(): self.gridgen_options[key] = gridgen_options[key] x = list(x); y = list(y) assert len(x)==len(y), 'arrays must be equal length' if beta is None: self.beta = [0 for xi in x] else: assert len(x)==len(beta), 'beta must have same length as x and y' self.beta = list(beta) self._line = Line2D(x, y, animated=True, ls='-', color='k', alpha=0.5, lw=1) self._ax.add_line(self._line) self._canvas = self._line.figure.canvas self._poly = Polygon(self.verts, alpha=0.1, fc='k', animated=True) self._ax.add_patch(self._poly) # Link in the lines that will show the beta values # pline for positive turns, mline for negative (minus) turns # otherwize zline (zero) for straight sections self._pline = Line2D([], [], marker='^', ms=12, mfc='g',\ animated=True, lw=0) self._mline = Line2D([], [], marker='v', ms=12, mfc='r',\ animated=True, lw=0) self._zline = Line2D([], [], marker='o', mfc='k', animated=True, lw=0) self._sline = Line2D([], [], marker='s', mfc='k', animated=True, lw=0) self._update_beta_lines() self._ax.add_artist(self._pline) self._ax.add_artist(self._mline) self._ax.add_artist(self._zline) self._ax.add_artist(self._sline) # get the canvas and connect the callback events cid = self._poly.add_callback(self._poly_changed) self._ind = None # the active vert self._canvas.mpl_connect('draw_event', self._draw_callback) self._canvas.mpl_connect('button_press_event',\ self._button_press_callback) self._canvas.mpl_connect('key_press_event', self._key_press_callback) self._canvas.mpl_connect('button_release_event',\ self._button_release_callback) self._canvas.mpl_connect('motion_notify_event',\ self._motion_notify_callback) def save_bry(self, bry_file='bry.pickle'): f = open(bry_file, 'wb') bry_dict = {'x': self.x, 'y': self.y, 'beta': self.beta} cPickle.dump(bry_dict, f, protocol=-1) f.close() def load_bry(self, bry_file='bry.pickle'): bry_dict = np.load(bry_file) x = bry_dict['x'] y = bry_dict['y'] self._line.set_data(x, y) self.beta = bry_dict['beta'] if hasattr(self, '_poly'): self._poly.xy = zip(x, y) self._update_beta_lines() self._draw_callback(None) self._canvas.draw() def save_grid(self, grid_file='grid.pickle'): f = open(grid_file, 'wb') cPickle.dump(self.grd, f, protocol=-1) f.close() def _get_verts(self): return zip(self.x, self.y) verts = property(_get_verts) def get_xdata(self): return self._line.get_xdata() x = property(get_xdata) def get_ydata(self): return self._line.get_ydata() y = property(get_ydata)