From 906ec537d81f9573fa44f30d7f2051f37e77c93f Mon Sep 17 00:00:00 2001
From: marco cammarata <marcocamma@gmail.com>
Date: Sun, 5 Mar 2017 16:40:39 +0100
Subject: [PATCH] improved masking, now one can add or subtract with GUI;
 implemented also polygon

---
 xray/mask.py | 176 +++++++++++++++++++++++++++++++++------------------
 1 file changed, 115 insertions(+), 61 deletions(-)

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) & (X<x2) & ( Y>y1) & (Y<y2)
+
+def _circleToMask(X,Y,vertices):
+    c,p = vertices
+    r = np.sqrt( (p[0]-c[0])**2 + (p[1]-c[1])**2 )
+    d = np.sqrt((X-c[0])**2+(Y-c[1])**2)
+    return d<r
+
+def _polygonToMask(X,Y,vertices):
+    points = np.vstack((X.flatten(),Y.flatten())).T
+    path = Path(vertices)
+    grid = path.contains_points(points)
+    return grid.reshape(X.shape)
 
 class MyMask(object):
   def __init__(self,img=None):
     self.comp = []
     self.img  = img
     self.mask = None
+    self._cache = None
+
+  def _define_component(self,operation,geometry,*vertices):
+    #print("define_comp",type(vertices),vertices)
+    if geometry == 'circle' and len(vertices) == 3:
+      xcen,ycen,radius = vertices
+      vertices = ( (xcen,ycen), (xcen+radius,ycen) )
+    if geometry == 'rectangle' and len(vertices) == 4:
+      vertices = ( (vertices[0],vertices[1]),(vertices[2],vertices[3]) )
+    # make sure vertices tuples
+    if isinstance(vertices,list):
+      vertices = [ (v[0],v[1]) for v in vertices ]
+      vertices = tuple(vertices)
+    a = dict( vertices = None )
+    self.comp.append( maskComponent(operation=operation,vertices=vertices,geometry=geometry) )
+    
   
-  def addCircle(self,xcen,ycen,radius):
-    self.comp.append( ["add","circle", [xcen,ycen,radius] ] )
-  def subtractCircle(self,xcen,ycen,radius):
-    self.comp.append( ["subtract","circle", [xcen,ycen,radius] ] )
-  def addRectangle(self,x1,y1,x2,y2):
-    if x1>x2: 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(d<r)
-        #raw_input()
-        m.append( d<r )
-      if kind == "rectangle":
-        (x1,y1,x2,y2) = pars
-        temp = (X>x1) & (X<x2) & ( Y>y1) & (Y<y2)
-        m.append( temp )
+    if shape is None and self.img is not None: shape = self.img.shape
+    if shape is None and self.img is     None: shape = self._cache['shape']
+    
+    if self._cache is None: self._cache = dict( shape = shape )
+    # reset cache if shape does not match
+    if shape != self._cache['shape']:
+      self._cache = dict( shape = shape )
+    X,Y = np.meshgrid ( range(shape[1]),range(shape[0]) )
+    for component in self.comp:
+      if component not in self._cache:
+        if component.geometry == 'circle':
+          mask = _circleToMask( X,Y,component.vertices )
+        elif component.geometry == 'rectangle':
+          mask = _rectangleToMask( X,Y,component.vertices )
+        elif component.geometry == 'polygon':
+          mask = _polygonToMask( X,Y,component.vertices )
+        else:
+          raise ValueError("Mask type %s not recongnized"%component.geometry)
+        self._cache[component] = mask
     mask = np.zeros(shape,dtype=np.bool)
-    for i in range(len(m)):
-      whattodo = self.comp[i][0]
-      if (whattodo == "add"):
-        mask[m[i]] = True
+    for comp in self.comp:
+      m = self._cache[ comp ]
+      if (comp.operation == "add"):
+        mask[m] = True
       else:
-        mask[m[i]] = False
+        mask[m] = False
     self.mask = mask
     return mask
 
@@ -70,25 +108,19 @@ class MyMask(object):
     i=fabio.edfimage.edfimage(mask.astype(np.uint8)); # edf does not support bool
     i.save(fname)
 
-def test():
-  mask = MyMask()
-  mask.addCircle(400,300,250)
-  mask.subtractCircle(400,300,150)
-  mask.addRectangle(350,250,1500,700)
-  mask.show()
-  return mask
-
 def snap(point,shape,snapRange=20):
   snapped = list(point)
   if snapped[0] < snapRange: snapped[0] = 0
-  if snapped[0] > 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()