Draw the borders of a binary Numpy array with Matplotlib - python-3.x

Question

I'm using image-segmentation on some images, and sometimes it would be nice to be able to plot the borders of the segments.
I have a 2D NumPy array that I plot with Matplotlib, and the closest I've gotten, is using contour-plotting.
This makes corners in the array, but is otherwise perfect.
Can Matplotlib's contour-function be made to only plot vertical/horizontal lines, or is there some other way to do this?
An example can be seen here:
import matplotlib.pyplot as plt
import numpy as np
array = np.zeros((20, 20))
array[4:7, 3:8] = 1
array[4:7, 12:15] = 1
array[7:15, 7:15] = 1
array[12:14, 13:14] = 0
plt.imshow(array, cmap='binary')
plt.contour(array, levels=[0.5], colors='g')
plt.show()

Answer 1

I wrote some functions to achieve this some time ago, but I would be glad to figure out how it can be done quicker.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
def get_all_edges(bool_img):
"""
Get a list of all edges (where the value changes from True to False) in the 2D boolean image.
The returned array edges has he dimension (n, 2, 2).
Edge i connects the pixels edges[i, 0, :] and edges[i, 1, :].
Note that the indices of a pixel also denote the coordinates of its lower left corner.
"""
edges = []
ii, jj = np.nonzero(bool_img)
for i, j in zip(ii, jj):
# North
if j == bool_img.shape[1]-1 or not bool_img[i, j+1]:
edges.append(np.array([[i, j+1],
[i+1, j+1]]))
# East
if i == bool_img.shape[0]-1 or not bool_img[i+1, j]:
edges.append(np.array([[i+1, j],
[i+1, j+1]]))
# South
if j == 0 or not bool_img[i, j-1]:
edges.append(np.array([[i, j],
[i+1, j]]))
# West
if i == 0 or not bool_img[i-1, j]:
edges.append(np.array([[i, j],
[i, j+1]]))
if not edges:
return np.zeros((0, 2, 2))
else:
return np.array(edges)
def close_loop_edges(edges):
"""
Combine thee edges defined by 'get_all_edges' to closed loops around objects.
If there are multiple disconnected objects a list of closed loops is returned.
Note that it's expected that all the edges are part of exactly one loop (but not necessarily the same one).
"""
loop_list = []
while edges.size != 0:
loop = [edges[0, 0], edges[0, 1]] # Start with first edge
edges = np.delete(edges, 0, axis=0)
while edges.size != 0:
# Get next edge (=edge with common node)
ij = np.nonzero((edges == loop[-1]).all(axis=2))
if ij[0].size > 0:
i = ij[0][0]
j = ij[1][0]
else:
loop.append(loop[0])
# Uncomment to to make the start of the loop invisible when plotting
# loop.append(loop[1])
break
loop.append(edges[i, (j + 1) % 2, :])
edges = np.delete(edges, i, axis=0)
loop_list.append(np.array(loop))
return loop_list
def plot_outlines(bool_img, ax=None, **kwargs):
if ax is None:
ax = plt.gca()
edges = get_all_edges(bool_img=bool_img)
edges = edges - 0.5 # convert indices to coordinates; TODO adjust according to image extent
outlines = close_loop_edges(edges=edges)
cl = LineCollection(outlines, **kwargs)
ax.add_collection(cl)
array = np.zeros((20, 20))
array[4:7, 3:8] = 1
array[4:7, 12:15] = 1
array[7:15, 7:15] = 1
array[12:14, 13:14] = 0
plt.figure()
plt.imshow(array, cmap='binary')
plot_outlines(array.T, lw=5, color='r')