python - 仅在 mask 区域计算梯度

Question

我有一个非常大的数组，其中只有几个感兴趣的小区域。我需要计算此数组的梯度，但出于性能原因，我需要将此计算限制在这些感兴趣的区域。

我不能做这样的事情:

phi_grad0[mask] = np.gradient(phi[mask], axis=0)

由于奇特的索引工作方式，phi[mask] 只是变成了屏蔽像素的一维数组，丢失了空间信息并使梯度计算毫无值(value)。

np.gradient 确实可以处理 np.ma.masked_array，但性能要差一个数量级:

import numpy as np
from timeit_context import timeit_context

phi = np.random.randint(low=-100, high=100, size=[100, 100])
phi_mask = np.random.randint(low=0, high=2, size=phi.shape, dtype=np.bool)

with timeit_context('full array'):
    for i2 in range(1000):
        phi_masked_grad1 = np.gradient(phi)

with timeit_context('masked_array'):
    phi_masked = np.ma.masked_array(phi, ~phi_mask)
    for i1 in range(1000):
        phi_masked_grad2 = np.gradient(phi_masked)

这会产生以下输出:

[full array] finished in 143 ms
[masked_array] finished in 1961 ms

我认为这是因为在 masked_array 上运行的操作未矢量化，但我不确定。

有没有办法限制np.gradient以获得更好的性能？

这个 timeit_context 是一个像这样工作的方便的计时器，如果有人感兴趣的话:

from contextlib import contextmanager
import time

@contextmanager
def timeit_context(name):
    """
    Use it to time a specific code snippet
    Usage: 'with timeit_context('Testcase1'):'
    :param name: Name of the context
    """
    start_time = time.time()
    yield
    elapsed_time = time.time() - start_time
    print('[{}] finished in {} ms'.format(name, int(elapsed_time * 1000)))

Answer 1

不完全是一个答案，但这是我设法针对我的情况拼凑起来的，效果很好:

我得到条件为真的像素的一维索引(例如，在这种情况下，条件为 < 5):

def get_indices_1d(image, band_thickness):
    return np.where(image.reshape(-1) < 5)[0]

这为我提供了一个包含这些索引的一维数组。

然后我以不同的方式手动计算这些位置的梯度:

def gradient_at_points1(image, indices_1d):
    width = image.shape[1]
    size = image.size

    # Using this instead of ravel() is more likely to produce a view instead of a copy
    raveled_image = image.reshape(-1)

    res_x = 0.5 * (raveled_image[(indices_1d + 1) % size] - raveled_image[(indices_1d - 1) % size])
    res_y = 0.5 * (raveled_image[(indices_1d + width) % size] - raveled_image[(indices_1d - width) % size])

    return [res_y, res_x]


def gradient_at_points2(image, indices_1d):
    indices_2d = np.unravel_index(indices_1d, dims=image.shape)

    # Even without doing the actual deltas this is already slower, and we'll have to check boundary conditions, etc
    res_x = 0.5 * (image[indices_2d] - image[indices_2d])
    res_y = 0.5 * (image[indices_2d] - image[indices_2d])

    return [res_y, res_x]


def gradient_at_points3(image, indices_1d):
    width = image.shape[1]

    raveled_image = image.reshape(-1)

    res_x = 0.5 * (raveled_image.take(indices_1d + 1, mode='wrap') - raveled_image.take(indices_1d - 1, mode='wrap'))
    res_y = 0.5 * (raveled_image.take(indices_1d + width, mode='wrap') - raveled_image.take(indices_1d - width, mode='wrap'))

    return [res_y, res_x]


def gradient_at_points4(image, indices_1d):
    width = image.shape[1]

    raveled_image = image.ravel()

    res_x = 0.5 * (raveled_image.take(indices_1d + 1, mode='wrap') - raveled_image.take(indices_1d - 1, mode='wrap'))
    res_y = 0.5 * (raveled_image.take(indices_1d + width, mode='wrap') - raveled_image.take(indices_1d - width, mode='wrap'))

    return [res_y, res_x]

我的测试数组是这样的:

a = np.random.randint(-10, 10, size=[512, 512])

# Force edges to not pass the condition
a[:, 0] = 99
a[:, -1] = 99
a[0, :] = 99
a[-1, :] = 99

indices = get_indices_1d(a, 5)

mask = a < 5

然后我可以运行这些测试:

with timeit_context('full gradient'):
    for i in range(100):
        grad1 = np.gradient(a)

with timeit_context('With masked_array'):
    for im in range(100):
        ma = np.ma.masked_array(a, mask)
        grad6 = np.gradient(ma)

with timeit_context('gradient at points 1'):
    for i1 in range(100):
        grad2 = gradient_at_points1(image=a, indices_1d=indices)

with timeit_context('gradient at points 2'):
    for i2 in range(100):
        grad3 = gradient_at_points2(image=a, indices_1d=indices)

with timeit_context('gradient at points 3'):
    for i3 in range(100):
        grad4 = gradient_at_points3(image=a, indices_1d=indices)

with timeit_context('gradient at points 4'):
    for i4 in range(100):
        grad5 = gradient_at_points4(image=a, indices_1d=indices)

结果如下:

[full gradient] finished in 576 ms
[With masked_array] finished in 3455 ms
[gradient at points 1] finished in 421 ms
[gradient at points 2] finished in 451 ms
[gradient at points 3] finished in 112 ms
[gradient at points 4] finished in 102 ms

如您所见，方法 4 是迄今为止最好的方法(不过不要太在意它消耗多少内存)。

这可能只是因为我的二维数组相对较小 (512x512)。也许对于更大的阵列，这不是真的。

另一个警告是 ndarray.take(indices, mode='wrap')将在图像边缘周围做一些奇怪的事情(一行将“循环”到下一行，等等)以保持良好的性能，因此如果边缘对您的应用程序很重要，您可能希望在边缘周围填充 1 个像素的输入数组.

还是 super 有趣多慢啊masked_array是。拉动构造函数ma = np.ma.masked_array(a, mask)循环外不影响自 masked_array 以来的时间本身只是保留对数组及其掩码的引用