comparison of dropletizing algorithms

dropletize_comparison.py

import numpy as np
import time
import matplotlib.pyplot as plt

g_fname = "beams_f16.npy"

_colors = ["#a6611a","#dfc27d","#80cdc1","#018571"]; # from colorbrewer2.prg

def maskEdges(img,edge=5,makeCopy=True):
  if edge is None or edge == 0: return img
  if makeCopy: img = img.copy()
  img[:edge,:]  = 0
  img[-edge:,:] = 0
  img[:,:edge]  = 0
  img[:,-edge:] = 0
  return img

def dropletImage_scipy(img,threshold=10, structure=np.ones( (3,3) ),edgeMask=5):
  """ threshold is very important !! """
  from scipy.ndimage import measurements
  img = maskEdges(img)
  label,nDrops = measurements.label(img>threshold,structure=structure)
  adu          = measurements.sum(img,label,index = range(1,nDrops+1) )
  cm           = measurements.center_of_mass(img,label,index=range(1,nDrops+1))
  drops = np.zeros( nDrops, dtype={'names': ["adu","x","y"],'formats' : [np.float16,np.float16,np.float16]})
  cm=np.asarray(cm)
  drops["x"]   = cm[:,0]
  drops["y"]   = cm[:,1]
  drops["adu"] = np.asarray(adu)
  return drops

def dropletImage_skimage(img,threshold=10, connectivity=2,edgeMask=5):
  from skimage import measure
  img = maskEdges(img)
  label = measure.label(img>threshold,connectivity=connectivity)
  props = measure.regionprops(label,intensity_image=img)
  adu   = [p.weighted_moments[0,0] for p in props]
  cm    = [p.weighted_centroid for p in props]
  drops = np.zeros( len(adu), dtype={'names': ["adu","x","y"],'formats' : [np.float16,np.float16,np.float16]})
  cm=np.asarray(cm)
  drops["x"]   = cm[:,0]
  drops["y"]   = cm[:,1]
  drops["adu"] = np.asarray(adu)
  return drops

def dropletImage_trackpy(img,diameter=5,edgeMask=5):
  import trackpy as tp
  img = maskEdges(img)
  f = tp.locate(img, diameter)
  drops = np.zeros( f.shape[0], dtype={'names': ["adu","x","y"],'formats' : [np.float16,np.float16,np.float16]})
  drops["x"]   = f["x"]
  drops["y"]   = f["y"]
  drops["adu"] = f["mass"]
  return drops

def dropletImage_skbeam(img,threshold=10,edgeMask=5):
  from skbeam.core.accumulators.droplet import dropletfind, dropletanal
  img = maskEdges(img)
  nDrops,label = dropletfind( (img>threshold).astype(int) )
  # the +0.5 offset is needed otherwise rounding will shift
  # the resulting ADU of the droplet
  npix, xcen, ycen, adus, idlist=dropletanal((img+0.5).astype(int),label,nDrops)
  drops = np.zeros( nDrops, dtype={'names': ["adu","x","y"],'formats' : [np.float16,np.float16,np.float16]})
  drops["x"]   = xcen
  drops["y"]   = ycen
  drops["adu"] = adus
  return drops
 
def testMethod(func,fname=g_fname,nImg=None):
  imgs = np.load(fname).astype(float)
  if nImg is not None: imgs=imgs[:nImg]
  t0 = time.time()
  drops = [func(img) for img in imgs]
  dt = time.time()-t0
  return np.concatenate(drops),dt
  
def doTest(nImg=5):
  fig,ax = plt.subplots(1,2)
  methods = (dropletImage_scipy,
             dropletImage_skimage,
             dropletImage_trackpy,
             dropletImage_skbeam) 
  for i,f in enumerate(methods):
    try:
      res,dt=testMethod(f,nImg=nImg)
      dt = dt/nImg*1e3
      label = "%s (%.0f ms)" % (f.__name__.split("_")[1],dt)
      kw = dict( histtype = 'step', color = _colors[i], alpha = 0.5,linewidth=3)
      ax[0].hist(res["adu"],np.arange(60,200),log=False,**kw)
      ax[1].hist(res["adu"],np.arange(0,1000),log=True,label=label,**kw)
    except ImportError:
      print("Can't run %s"%f.__name__)
  ax[0].set_ylabel( "Num of drops in %d images" % nImg)
  ax[0].set_title("One photon peak")
  ax[0].set_xlabel("ADU")
  ax[1].set_xlabel("ADU")
  plt.legend()

if __name__ == "__main__": doTest(10)