#!/usr/bin/env python
# coding: utf-8

# > This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), 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
get_ipython().run_line_magic('matplotlib', 'inline')


# 2. 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)


# 3. 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)


# 4. 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


# 5. 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);


# 6. 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);


# 7. 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)


# 8. 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](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages).