diff --git a/xray/mask.py b/xray/mask.py index ea51c60..f2c8a1e 100644 --- a/xray/mask.py +++ b/xray/mask.py @@ -5,53 +5,91 @@ import logging log = logging.getLogger(__name__) import os +import collections import numpy as np import matplotlib.pyplot as plt +from matplotlib.path import Path + +maskComponent = collections.namedtuple('maskComponent',['operation','geometry','vertices']) + +def _rectangleToMask(X,Y,vertices): + ( (x1,y1), (x2,y2) ) = vertices + if x1>x2: x1,x2=x2,x1 + if y1>y2: y1,y2=y2,y1 + return (X>x1) & (Xy1) & (Yx2: x1,x2=x2,x1 - if y1>y2: y1,y2=y2,y1 - self.comp.append( ["add","rectangle", [x1,y1,x2,y2] ]) - def subtractRectangle(self,x1,y1,x2,y2): - if x1>x2: x1,x2=x2,x1 - if y1>y2: y1,y2=y2,y1 - self.comp.append( ["subtract","rectangle", [x1,y1,x2,y2] ]) + def addCircle(self,*vertices): self._define_component( 'add', 'circle', *vertices ) + def subtractCircle(self,*vertices): self._define_component( 'subtract', 'circle', *vertices ) + + def addRectangle(self,*vertices): self._define_component( 'add','rectangle', *vertices) + def subtractRectangle(self,*vertices): self._define_component( 'subtract','rectangle',*vertices) + + + def addPolygon(self,*vertices): self._define_component( 'add','polygon',*vertices) + def subtractPolygon(self,*vertices): self._define_component( 'subtract','polygon',*vertices) + def getMask(self,shape=None): - if shape is None: shape = self.img.shape - m = [] - X,Y = np.meshgrid ( range(shape[0]),range(shape[1]) ) - for o in self.comp: - whattodo = o[0] - kind=o[1] - pars=o[2] - if kind == "circle": - (xc,yc,r) = pars - d = np.sqrt((X-xc)**2+(Y-yc)**2) - #plt.imshow(dx1) & (Xy1) & (Y shape[0]-snapRange: snapped[0] = shape[0] + if snapped[0] > shape[1]-snapRange: snapped[0] = shape[1] if snapped[1] < snapRange: snapped[1] = 0 - if snapped[1] > shape[1]-snapRange: snapped[1] = shape[1] - return snapped + if snapped[1] > shape[0]-snapRange: snapped[1] = shape[0] + return tuple(snapped) -def getPoint(shape,snapRange): - c = plt.ginput()[0] - c = snap(c,shape,snapRange=snapRange) +def getPoints(N=1,shape=(100,100),snapRange=0): + if N<1: print('Right click cancels last point, middle click ends the polygon') + c = plt.ginput(N) + c = [ snap(point,shape,snapRange=snapRange) for point in c ] + if len(c) == 1: c = c[0] return c def makeMaskGui(img,snapRange=60): @@ -100,20 +132,32 @@ def makeMaskGui(img,snapRange=60): plt.clim(np.percentile(img,(2,98))) plt.imshow(mask.getMatplotlibMask()) plt.pause(0.01) - ans = input("What's next c/r/done? ") + ans = input("What's next p/P/c/C/r/R/done? (capitals = subtract)") if ans == "c": - print("Adding circle, click on center") - c = getPoint(img.shape,snapRange) - print("Adding circle, click on another point to define radius") - p = getPoint(img.shape,snapRange) - r = np.sqrt( (p[0]-c[0])**2 + (p[1]-c[1])**2 ) - mask.addCircle(c[0],c[1],r) + print("Adding circle, click on center then another point to define radius") + vertices = getPoints(N=2,shape=img.shape,snapRange=snapRange) + mask.addCircle(*vertices) + if ans == "C": + print("Subtracting circle, click on center then another point to define radius") + vertices = getPoints(N=2,shape=img.shape,snapRange=snapRange) + mask.subtractCircle(*vertices) if ans == "r": - print("Adding rectangle, click on one corner") - c1 = getPoint(img.shape,snapRange) - print("Adding rectangle, click on opposite corner") - c2 = getPoint(img.shape,snapRange) - mask.addRectangle(c1[0],c1[1],c2[0],c2[1]) + print("Adding rectangle, click on one corner and then on the opposite one") + vertices = getPoints(N=2,shape=img.shape,snapRange=snapRange) + mask.addRectangle(*vertices) + if ans == "R": + print("Subtracting rectangle, click on one corner and then on the opposite one") + vertices = getPoints(N=2,shape=img.shape,snapRange=snapRange) + mask.subtractRectangle(*vertices) + if ans == 'p': + print("Adding polygon") + vertices = getPoints(N=-1,shape=img.shape,snapRange=snapRange) + mask.addPolygon(*vertices) + if ans == 'P': + print("Subtracting polygon") + vertices = getPoints(N=-1,shape=img.shape,snapRange=snapRange) + mask.subtractPolygon(*vertices) + plt.imshow(mask.getMatplotlibMask()) plt.pause(0.01) fname = input("Enter a valid filename (ext .edf or .npy) if you want to save the mask (empty otherwise)") @@ -146,10 +190,20 @@ def maskCenterLines(width,shape): w0 = int(width[0]/2) w1 = int(width[1]/2) mask[ c0-w0:c0+w0 , : ] = True - mask[ : , c1-w1:c1+w1 ] = True + mask[ : , c1-w1:c1+w1 ] = True return mask +def test(shape=(1000,2000)): + mask = MyMask() + mask.addCircle(400,300,250) + mask.subtractCircle(400,300,150) + mask.addRectangle(350,250,1500,700) + plt.imshow( mask.getMask(shape) ) + return mask + + + if __name__ == "__main__": test() plt.show()