DZO Assignment: Random Walker-Based Segmentation

In this lab, we will deal with image segmentation. This will involve computing the weights between neighboring pixels; the weights will reflect how similar the pixels are. Then the segmentation will be seeded (= some pixels will be marked as foreground/background) and finally the segmentation will be computed.

We will start with a toy image, making sure all our implementations are correct. Then we will shift to a realistic scenario with a microscopic image.

%load_ext autoreload
%autoreload 2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

from segmentation import *

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

# Let us create a toy image consisting of two areas, 
# one dark and the other one bright.
im = np.zeros((4, 4))
im[:, 2:] = 1.0
plt.imshow(im, cmap='gray'); plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f2bc8b88370>

# Let us create the mask which defines which points are foreground (mask == 1) 
# and which are background (mask==2). 
mask = np.zeros((4, 4), dtype=int)
# set bottom right corner to FG:
mask[-1, -1] = 1
# set top left corner to BG:
mask[0,0] = 2
plt.imshow(mask); plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f2bc8a47af0>

# define the sigma for computing the weights between pixels:
sigma = 0.5
# and prepare the variables needed for further computations: 
x, weights_r, weights_l, weights_t, weights_b, mask_padded, diff_v, diff_h = seg_prepare(im, mask, sigma)

Let us have a look at what we've prepared.

show_tuple((x, mask_padded), ('initial x', 'padded mask') )

These two arrays show what will is true about every output array from the seg_prepare function (except diff_v and diff_h which are included just for debug purposes): All the arrays are padded by a border of size 1 so that further computations do not have to deal with boundary conditions using if statements or similar.

The initial solution for forground probability (initial x) is generated just randomly (np.random.rand). The values in the padded region of x are set to -1: this is an arbitrary choice and it helps to debug the code. If any update of x produces negative value in the valid (non-padded) region of x, there must be a bug in the implementation (because probability can't be negative.)

The weights_X variables store the computed weights between the central pixel and its left, right, top and bottom neighbors, respectively:

show_tuple((weights_r, weights_l, weights_t, weights_b), ['weights_r', 'weights_l', 'weights_t', 'weights_b'])

You can see an example use of these weights in the computation of Random Walker criterion in the provided function compute_crit.

s = compute_crit(x,  weights_r, weights_l, weights_t, weights_b)
print(f'Initial criterion value: {s}')

Initial criterion value: 3.512200138919708

Now we have everything ready for running an iterative optimization on x. Implement the function GS_iteration based on Gauss Seidel iterations. The function will produce an update of x.

Run the function for a fixed number of iterations and store the criterion values:

criterion = []
iters = 30
for it in range(iters):
    x = GS_iteration(x, mask_padded, weights_r, weights_l, weights_t, weights_b)
    s = compute_crit(x,  weights_r, weights_l, weights_t, weights_b)
    criterion.append(s)

plt.plot(criterion)

[<matplotlib.lines.Line2D at 0x7f2bc85a08e0>]

Good, the optimization has converted. Let us have a look at the solution:

plt.imshow(x[1:-1, 1:-1]); plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f2bc840b370>

Let us threshold the solution in the middle (0.5) to get a hard segmentation:

plt.imshow(x[1:-1, 1:-1]>0.5); plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f2bc84d3100>

Now let us try to the same for a microscopic image.

im = Image.open('microscopic.png').convert('L')
im = np.array(im).astype(float)/255.0
plt.imshow(im, cmap='gray'); plt.colorbar() 

<matplotlib.colorbar.Colorbar at 0x7f2bc8f7db10>

Let us seed using the very high and very low values:

mask = np.zeros_like(im, dtype=int)
# create seeds
mask[im>0.98] = 1
mask[im<0.02] = 2 

plt.imshow(mask); plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f2bc861c1c0>

x, weights_r, weights_l, weights_t, weights_b, mask_padded, diff_v, diff_h = seg_prepare(im, mask, sigma)

# show initial x: 
plt.imshow(x[1:-1, 1:-1]); plt.colorbar(); plt.title('initial x')

Text(0.5, 1.0, 'initial x')

# run the optimization:
criterion = []
for k in range(30):
    x = GS_iteration(x, mask_padded, weights_r, weights_l, weights_t, weights_b)
    s = compute_crit(x,  weights_r, weights_l, weights_t, weights_b)
    criterion.append(s)

plt.plot(criterion)

[<matplotlib.lines.Line2D at 0x7f2bc8857e20>]

plt.imshow(x[1:-1, 1:-1]>0.5); plt.title('Segmentation')

Text(0.5, 1.0, 'Segmentation')

And finally, compare this to threshold-based method:

plt.imshow(im[1:-1, 1:-1]>0.5); plt.title('Thresholded image intensity')

Text(0.5, 1.0, 'Thresholded image intensity')