Векторизованное извлечение подматриц в numpy -матрице - PullRequest
Новая
       75

Векторизованное извлечение подматриц в numpy -матрице

4 голосов
/ 02 мая 2020

Моя цель - реализовать медианный фильтр, который является функцией, которая заменяет каждый пиксель в (в основном) 2d-массиве медианой окружающих его пикселей. Его можно использовать для шумоподавления изображений.

Моя реализация извлекает подматрицы из исходной матрицы, которые содержат сам пиксель и его соседей. Это извлечение в настоящее время выполняется для for-l oop, и, как вы можете себе представить, for-l oop отнимает примерно 95% времени выполнения.

Вот моя текущая реализация: 

def median_blur(img, fsize=3):
    img_intermediate = np.zeros((img.shape[0] + 2, img.shape[1] + 2), np.uint8)  # First intermediate img, used for padding original image
    img_intermediate[1:img.shape[0]+1, 1:img.shape[1]+1] = img
    img_result = np.empty((*img.shape, fsize, fsize), np.uint8)  # Will contain result, first receives kernel-submatrices

    # Extract submatrices from image
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            img_result[i, j] = img_intermediate[i:i+fsize, j:j+fsize]

    img_result = img_result.reshape(*img.shape, fsize**2)  # Reshape Result-Matrix
    img_result = np.median(img_result, axis=2)  # Calculate median in one go
    return img_result.astype('uint8')

Как я могу извлечь эти подматрицы с помощью векторизованной операции?

В качестве бонуса, если кто-то с опытом работы в компьютерном зрении читает это: есть ли лучший подход реализации медианного фильтра, отличного от применения его к промежуточной матрице с нулями?

Большое спасибо.

1 Ответ

3 голосов
/ 02 мая 2020

Вот векторизованное решение. Однако вы можете найти более быстрое решение, обращая внимание на порядок памяти массива изображений: 

from numpy.lib.stride_tricks import as_strided

img_padded = np.pad(img, 1, mode='constant')
sub_shape = (fsize, fsize)
view_shape = tuple(np.subtract(img_padded.shape, sub_shape) + 1) + sub_shape
sub_img = as_strided(img_padded, view_shape, img_padded.strides * 2)
sub_img = sub_img.reshape(img.shape + (fsize**2,))
result = np.median(sub_img, axis=2).astype(int)

Используйте функцию pad для заполнения сначала, а затем as_strided для получения sub_matrices и, наконец, применения median к вашим шагам.

ОБНОВЛЕНИЕ : использование view_as_windows, предложенного @Divakar в комментариях: 

from skimage.util.shape import view_as_windows

img_padded = np.pad(img, 1, mode='constant')
sub_shape = (fsize, fsize)
view_shape = view_as_windows(img_padded, sub_shape, 1)
sub_img = view_shape.reshape(img.shape + (fsize**2,))
result = np.median(sub_img, axis=2).astype(int)

view_as_windows также обеспечивает подобласти подматрицы к шагам.

образец изображения и вывод: 

img: 
[[ 1  2  3  4  5  6]
 [ 7  8  9 10 11 12]
 [13 14 15 16 17 18]
 [19 20 21 22 23 24]]

median_filtered: 
[[ 0  2  3  4  5  0]
 [ 2  8  9 10 11  6]
 [ 8 14 15 16 17 12]
 [ 0 14 15 16 17  0]]
...