This is one of the 100 recipes of the IPython Cookbook, the definitive guide to high-performance scientific computing and data science in Python.

14.4. Computing connected components in an image

  1. Let's import the packages.
In [ ]:
import itertools
import numpy as np
import networkx as nx
import matplotlib.colors as col
import matplotlib.pyplot as plt
%matplotlib inline
  1. We create a $10 \times 10$ image where each pixel can take one of three possible labels (or colors).
In [ ]:
n = 10
In [ ]:
img = np.random.randint(size=(n, n), 
                        low=0, high=3)
  1. Now, we create the underlying 2D grid graph encoding the structure of the image. Each node is a pixel, and a node is connected to its nearest neighbors. NetworkX defines a function grid_2d_graph for generating this graph.
In [ ]:
g = nx.grid_2d_graph(n, n)
  1. Let's create two functions to display the image and the corresponding graph.
In [ ]:
def show_image(img, **kwargs):
    plt.imshow(img, origin='lower',interpolation='none',
               **kwargs);
    plt.axis('off');
In [ ]:
def show_graph(g, **kwargs):
    nx.draw(g,
            pos={(i, j): (j, i) for (i, j) in g.nodes()},
            node_color=[img[i, j] for (i, j) in g.nodes()],
            linewidths=1, edge_color='w', with_labels=False,
            node_size=30, **kwargs);
In [ ]:
cmap = plt.cm.Blues
  1. Here is the original image superimposed with the underlying graph.
In [ ]:
plt.figure(figsize=(5,5));
show_image(img, cmap=cmap, vmin=-1);
show_graph(g, cmap=cmap, vmin=-1);
  1. We are now going to find all the contiguous regions of the image in dark blue containing more than three pixels. The first step is to consider the subgraph corresponding to all dark blue pixels.
In [ ]:
g2 = g.subgraph(zip(*np.nonzero(img==2)))
In [ ]:
plt.figure(figsize=(5,5));
show_image(img, cmap=cmap, vmin=-1);
show_graph(g2, cmap=cmap, vmin=-1);
  1. We see that the requested contiguous regions correspond to the connected components containing more than three nodes in the subgraph. We can use the connected_components function of NetworkX to find those components.
In [ ]:
components = [np.array(comp)
              for comp in nx.connected_components(g2)
              if len(comp)>=3]
len(components)
  1. Finally, let's assign a new color to each of these components, and let's display the new image.
In [ ]:
# We copy the image, and assign a new label
# to each found component.
img_bis = img.copy()
for i, comp in enumerate(components):
    img_bis[comp[:,0], comp[:,1]] = i + 3
In [ ]:
# We create a new discrete color map extending
# the previous map with new colors.
colors = [cmap(.5), cmap(.75), cmap(1.), 
          '#f4f235', '#f4a535', '#f44b35']
cmap2 = col.ListedColormap(colors, 'indexed')
In [ ]:
plt.figure(figsize=(5,5));
show_image(img_bis, cmap=cmap2);

You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).

IPython Cookbook, by Cyrille Rossant, Packt Publishing, 2014 (500 pages).