Image Segmentation using Active Contours, Level Sets

Here is a new image segmentation Python code which is a port of this Matlab file. It uses level sets and mean curvature motion, and is able to segment the sample image after few iterations. We had to write our own bwdist which is a wrapper for distance_transform_edt wrapper so bwdist acted the same way as Matlab bwdist(), some conversions were necessary to make this work. We also changed the initialization of \phi, making it simpler and closer to Dr. Li's code we shared before.

Code

import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as signal
import scipy.ndimage as image
import time
from scipy import ndimage

def bwdist(a):
   """
   this is an intermediary function, 'a' has only True, False vals,
   so we convert them into 0, 1 values -- in reverse. True is 0,
   False is 1, distance_transform_edt wants it that way.
   """
   b = np.ones(a.shape)
   b[a==True] = 0.
   return ndimage.distance_transform_edt(b)

def gauss_kern():
   """ Returns a normalized 2D gauss kernel array for convolutions """
   h1 = 15
   h2 = 15
   x, y = np.mgrid[0:h2, 0:h1]
   x = x-h2/2
   y = y-h1/2
   sigma = 1.5
   g = np.exp( -( x**2 + y**2 ) / (2*sigma**2) );
   return g / g.sum()

Img = plt.imread("twoObj.bmp")
Img = Img[::-1]
g = gauss_kern()
Img_smooth = signal.convolve(Img,g,mode='same')
Iy,Ix=np.gradient(Img_smooth)
absGradI=np.sqrt(Ix**2+Iy**2);
rows, cols = Img.shape

# initial function phi - level set is a square 4 pixels
# away from borders on each side, in 3D it looks like an empty
# box
c0=4
w=4
nrow, ncol=Img.shape
phi=c0*np.ones((nrow,ncol))
phi[w+1:-w-1, w+1:-w-1]=-c0

# edge-stopping function
g = 1 / (1+absGradI**2)

# gradient of edge-stopping function
gy,gx = np.gradient(g)

# gradient descent step size
dt=.4

# number of iterations after which we reinitialize the surface
num_reinit=10

phiOld=np.zeros((rows,cols))

# number of iterations after which we reinitialize the surface
iter=0

plt.ion()

while np.sum(np.sum(np.abs(phi-phiOld))) != 0:
   # gradient of phi
   gradPhiY, gradPhiX = np.gradient(phi)   
   # magnitude of gradient of phi
   absGradPhi=np.sqrt(gradPhiX**2+gradPhiY**2)
   # normalized gradient of phi - eliminating singularities
   normGradPhiX=gradPhiX/(absGradPhi+(absGradPhi==0))
   normGradPhiY=gradPhiY/(absGradPhi+(absGradPhi==0))
  
   divYnormGradPhiX, divXnormGradPhiX=np.gradient(normGradPhiX)
   divYnormGradPhiY, divXnormGradPhiY=np.gradient(normGradPhiY)
                         
   # curvature is the divergence of normalized gradient of phi
   K = divXnormGradPhiX + divYnormGradPhiY   
   tmp1 = g * K * absGradPhi
   tmp2 = g * absGradPhi
   tmp3 = gx * gradPhiX + gy*gradPhiY
   dPhiBydT =tmp1 + tmp2 + tmp3
   phiOld=phi
   # level set evolution equation   
   phi = phi + ( dt * dPhiBydT )
   iter=iter+1
   if np.mod(iter,num_reinit)==0:
       # reinitialize the embedding function after num_reinit iterations
       phi=np.sign(phi)
       phi = (phi > 0) * (bwdist(phi <> 0))
      
   if np.mod(iter,10)==0:
       time.sleep(0.6)
       plt.imshow(Img, cmap='gray')
       plt.hold(True)
       CS = plt.contour(phi,0, colors='r')
       plt.draw()
       plt.hold(False)