Distance to Nearest 0 Value for every pixel in 2D Binary Array-CodePudding

Given a 2D binary array (only 1s or 0s), for the pixels that have a value of 1, I want to find the distance from that pixel to the nearest 0-valued pixel, and replace the 1-valued pixel with that distance. If you're familiar with the game Minesweeper, you can imagine that the 0-valued pixels are mines, and we want to assign proximity values to all other pixels. Exact precision is not necessarily important; Minesweeper will assign a pixel that is diagonal to a mine as valued as 1 (as opposed to sqrt(2)). So whatever is fastest/easiest. So for example:

[[1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1],
 [1, 1, 0, 1, 1],
 [1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1],
 ]

would become:

[[2, 2, 2, 2, 2],
 [2, 1, 1, 1, 2],
 [2, 1, 0, 1, 2],
 [2, 1, 1, 1, 2],
 [2, 2, 2, 2, 2],
 ]

A naive way would be to get the 2D indices of all 0-valued pixels and 1-valued pixels, calculate the euclidean distance matrix between the two, find the minimum value for each 1-valued pixel, then assign those values to the corresponding indices.

import numpy as np
# example sample
maze = np.ones((100, 100))
maze[40:60, 40:60] = 0

# find indices
mask_z = maze == 0
idx_z = np.nonzero(mask_z)
idx_nz = np.nonzero(~mask_z)
idx_z = np.stack(idx_z, axis=-1)
idx_nz = np.stack(idx_nz, axis=-1)

# calculate distance matrix
dist_matrix = np.linalg.norm(idx_nz.reshape(-1, 1, 2) - idx_z, axis=-1)
dists = np.min(dist_matrix, axis=-1)

# assign distances
mapped = np.zeros(floorplan.shape[:2], dtype=int)
mapped[idx_nz] = dists.squeeze()

However, the issue is that for large matrices (2k by 2k) I run out of memory and speed is pretty poor. Is there a better way to do this?

CodePudding user response：

You can either use scipy's distance_transform_cdt or opencv's distanceTransform.

A personal note, I just ran into a weird case where distanceTransform was buggy and had to use distance_transform_edt, so I'd recommend the latter.