Vectorialize getting the min of elements in a position relative to the current element with numpy

Question

I'm having a hard time vectorializing the following function:

def find_empty_square_sizes(matrix: np.ndarray):
    shape = matrix.shape
    res_matrix = np.ones(shape, dtype=np.int32)
    res_matrix[matrix != 0] = 0
    for i in range(1, shape[0]):
        for j in range(1, shape[1]):
            if matrix[i][j] == 1:
                res_matrix[i][j] = 0
            else:
                res_matrix[i][j] = np.min(np.ma.masked_array(res_matrix[i - 1:i + 1, j - 1:j + 1], mask=[[0, 0], [0, 1]])) + 1
    return res_matrix

The idea is to find the biggest square sub matrix of zeros inside a matrix. What the result, res_matrix, means is if the element on row i and column j was the bottom right corner of a sub matrix, how big that sub matrix could be while having only zeros.

Running for example the following code:

m = np.zeros((6, 6), dtype=np.int32)
m[2, 2] = 1
res = find_empty_square_sizes(m)
print(m)
print(res)

Yields the following results:

[[0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 1 0 0 0]
 [0 0 0 1 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]]
[[1 1 1 1 1 1]
 [1 2 2 2 2 2]
 [1 2 0 1 2 3]
 [1 2 1 0 1 2]
 [1 2 2 1 1 2]
 [1 2 3 2 2 2]]

I'm doing this for some image processing purposes and for HD images with over a thousand rows and columns this takes about a minute on my PC and I'd like to improve that. Any ideas on how to vectorialize the find_empty_square_sizes function or achieve a similar result through a more efficient method?

Answer 1

If anyone stumbles upon this with a similar problem, following Jerome Richard's advice I used Cython and the following code brought down the execution time from a minute to about 40 ms, quite a difference!

#cython: language_level=3

import cython
import numpy as np
cimport numpy as np


@cython.boundscheck(False)
@cython.wraparound(False)


def find_empty_square_sizes(matrix: np.ndarray) -> np.ndarray:
    res_matrix = np.ones((matrix.shape[0], matrix.shape[1]), dtype=np.int32)
    res_matrix[matrix != 0] = 0
    _find_empty_square_sizes(matrix, res_matrix, matrix.shape[0], matrix.shape[1])
    return res_matrix


cdef _find_empty_square_sizes(int[:, :] matrix, int[:, :] res_matrix, int x, int y):
    cdef int i
    cdef int j

    for i in range(1, x):
        for j in range(1, y):
            if matrix[i][j] == 1:
                res_matrix[i][j] = 0
            else:
                res_matrix[i][j] = _min(res_matrix[i - 1, j - 1], res_matrix[i, j - 1], res_matrix[i - 1, j]) + 1


cdef _min(int a, int b, int c):
    cdef int min = a
    if b < min:
        min = b
    if c < min:
        min = c
    return min